@echo off
setlocal EnableDelayedExpansion

set FQBN=m5stack:esp32:m5stack_atoms3:USBMode=default,CDCOnBoot=cdc,UploadMode=default,PartitionScheme=default_8MB
set SKETCH=%~dp0..\firmware\MacroPad
set BEFORE=%TEMP%\macropad_ports_before.txt
set AFTER=%TEMP%\macropad_ports_after.txt
set DIFF=%TEMP%\macropad_ports_diff.txt
set ESPTOOL=
for /f "delims=" %%E in ('powershell -NoProfile -Command "Get-ChildItem -Path \"$env:LOCALAPPDATA\Arduino15\packages\" -Recurse -Filter esptool.exe -ErrorAction SilentlyContinue | Sort-Object LastWriteTime -Descending | Select-Object -First 1 -ExpandProperty FullName"') do set ESPTOOL=%%E

if not defined ESPTOOL (
    echo ERROR: esptool.exe not found under %%LOCALAPPDATA%%\Arduino15\packages
    echo The Arduino M5Stack/ESP32 board package may not be installed.
    pause
    exit /b 1
)

set EXIT_CODE=0
set AUTO_PORT_FILE=%TEMP%\macropad_auto_port.txt

:: --lite only changes the on-screen instructions (the AtomS3 Lite has no
:: display to watch). The firmware itself is universal - same FQBN, same
:: binary - and detects the board at boot.
set LITE=
set EXPLICIT_PORT=
if /i "%~1"=="--lite" (
    set LITE=1
    set EXPLICIT_PORT=%2
) else (
    set EXPLICIT_PORT=%1
    if /i "%~2"=="--lite" set LITE=1
)

:: Power-user shortcut: explicit port skips the interactive flow.
if not "%EXPLICIT_PORT%"=="" (
    set PORT=%EXPLICIT_PORT%
    goto :compile
)

echo.
if defined LITE (
    echo ===============================================
    echo   ATOM S3 Lite MacroPad - Interactive Flasher
    echo ===============================================
    echo.
    echo  NOTE: the Lite has no screen. Status comes from
    echo  its RGB LED - after flashing, a short white pulse
    echo  means the firmware booted.
) else (
    echo ===============================================
    echo   ATOM S3 MacroPad - Interactive Flasher
    echo ===============================================
)
echo.

:: -----------------------------------------------
::  Auto path: if our firmware is already running on
::  a connected ATOMS3, ask it to drop into the ROM
::  bootloader and skip the hold-the-button dance.
::  Falls through to the manual flow on any failure.
:: -----------------------------------------------
echo Checking for a running MacroPad on USB...
echo.
if exist "%AUTO_PORT_FILE%" del "%AUTO_PORT_FILE%" >nul 2>&1
python "%~dp0auto_enter_bootloader.py" "%AUTO_PORT_FILE%"
set AUTO_RESULT=%ERRORLEVEL%
if %AUTO_RESULT%==0 if exist "%AUTO_PORT_FILE%" (
    set /p PORT=<"%AUTO_PORT_FILE%"
    if not "!PORT!"=="" (
        echo.
        echo Auto-detected ATOMS3 in download mode on !PORT!.
        echo Skipping manual reconnect steps.
        echo.
        goto :compile
    )
)
echo.
echo Auto-detection did not find a running MacroPad.
echo Falling back to manual reconnect flow.
echo.
echo This script will detect the device automatically by
echo watching which COM port appears when you plug the ATOMS3 in.
echo.

:: ---- Step 1: snapshot ports while device is unplugged ----
echo -----------------------------------------------
echo  STEP 1 of 4: Disconnect the ATOM S3
echo -----------------------------------------------
echo.
echo  Unplug the ATOM S3 from your computer if it's
echo  currently plugged in. If it's not plugged in,
echo  no action is needed.
echo.
pause

powershell -NoProfile -Command "$ports = @([System.IO.Ports.SerialPort]::GetPortNames() | Sort-Object -Unique); [System.IO.File]::WriteAllLines('%BEFORE%', $ports)"
if not exist "%BEFORE%" (
    echo Failed to enumerate COM ports. PowerShell may be restricted on this system.
    set EXIT_CODE=1
    goto :end
)

echo.
echo  Currently visible COM ports:
set _ANY=
for /f "usebackq delims=" %%L in ("%BEFORE%") do (
    echo    %%L
    set _ANY=1
)
if not defined _ANY echo    (none)
echo.

:: ---- Step 2: ask user to plug in + enter download mode ----
:reconnect_step
echo -----------------------------------------------
echo  STEP 2 of 4: Plug in and enter download mode
echo -----------------------------------------------
echo.
if defined LITE (
    echo  Plug the ATOM S3 Lite in, then press AND HOLD
    echo  the small side button until the green light
    echo  flashes. That puts it into download mode.
    echo.
    echo  ^(The Lite's main RGB LED stays dark in download
    echo  mode - the green flash is the small internal
    echo  light next to the USB port.^)
) else (
    echo  Plug the ATOM S3 in, then press AND HOLD the
    echo  small side button until the internal green
    echo  light flashes. That puts it into download mode.
)
echo.
echo  (If you release too early it'll boot normally.
echo  Just unplug, replug, and try again.)
echo.
pause

powershell -NoProfile -Command "$ports = @([System.IO.Ports.SerialPort]::GetPortNames() | Sort-Object -Unique); [System.IO.File]::WriteAllLines('%AFTER%', $ports)"
if not exist "%AFTER%" (
    echo Failed to enumerate COM ports.
    set EXIT_CODE=1
    goto :end
)

:: ---- Step 3: diff (ports in AFTER not in BEFORE) -> %DIFF% ----
powershell -NoProfile -Command "$b = @(Get-Content '%BEFORE%' -ErrorAction SilentlyContinue); $a = @(Get-Content '%AFTER%' -ErrorAction SilentlyContinue); $new = @($a | Where-Object { $b -notcontains $_ }); [System.IO.File]::WriteAllLines('%DIFF%', $new)"

set NEW_COUNT=0
set PORT=
for /f "usebackq delims=" %%P in ("%DIFF%") do (
    set /a NEW_COUNT+=1
    if not defined PORT set PORT=%%P
)

if !NEW_COUNT! GTR 1 (
    echo.
    echo  More than one new COM port appeared:
    for /f "usebackq delims=" %%L in ("%DIFF%") do echo    %%L
    echo.
    echo  Cannot determine which is the ATOM S3. Unplug
    echo  any other USB-serial devices and try again, or
    echo  run "upload.bat COMx" with the port specified.
    set EXIT_CODE=1
    goto :end
)

if !NEW_COUNT! EQU 0 (
    echo.
    echo  No new COM port detected.
    echo  Most likely you didn't hold the side button long
    echo  enough, or didn't reconnect the device.
    echo.
    set RETRY=
    set /p RETRY=  Try step 2 again? [Y/N]:
    if /i "!RETRY!"=="Y" goto :reconnect_step
    echo  Aborted.
    set EXIT_CODE=1
    goto :end
)

:: ---- Step 4: confirm + countdown ----
echo.
echo -----------------------------------------------
echo  STEP 3 of 4: Confirm
echo -----------------------------------------------
echo.
echo  Detected ATOM S3 in download mode on !PORT!.
echo.
set CONFIRM=
set /p CONFIRM=  Flash this device? [Y/N]:
if /i not "!CONFIRM!"=="Y" (
    echo  Aborted.
    set EXIT_CODE=1
    goto :end
)

echo.
echo -----------------------------------------------
echo  STEP 4 of 4: Flashing in 5 seconds...
echo -----------------------------------------------
echo  Press Ctrl+C now to abort.
echo.
for /l %%i in (5,-1,1) do (
    echo    %%i...
    ping -n 2 127.0.0.1 >nul
)
echo.

:compile
echo ===============================================
echo   Flashing MacroPad firmware
echo   Port: %PORT%
echo   FQBN: %FQBN%
echo ===============================================
echo.

echo -----------------------------------------------
echo  Step 1 of 3: Compiling firmware
echo -----------------------------------------------
arduino-cli compile --fqbn "%FQBN%" "%SKETCH%"
set RESULT=%ERRORLEVEL%
if not %RESULT%==0 (
    echo.
    echo COMPILE FAILED - cannot upload.
    set EXIT_CODE=%RESULT%
    goto :end
)
echo.
echo   Compile OK
echo.

echo -----------------------------------------------
echo  Step 2 of 3: Wiping device flash
echo -----------------------------------------------
echo  Full chip erase - clears firmware, NVS, and
echo  LittleFS so the device boots up clean.
echo.
"%ESPTOOL%" --chip esp32s3 --port %PORT% erase_flash
set RESULT=%ERRORLEVEL%
if not %RESULT%==0 (
    echo.
    echo ERASE FAILED.
    echo.
    echo If the device dropped out of download mode, hold the
    echo side button while plugging it in, then re-run.
    set EXIT_CODE=%RESULT%
    goto :end
)
echo.
echo   Wipe OK - device flash is now empty
echo.

echo -----------------------------------------------
echo  Step 3 of 3: Uploading new firmware
echo -----------------------------------------------
arduino-cli upload -v --fqbn "%FQBN%" --port %PORT% "%SKETCH%"
set RESULT=%ERRORLEVEL%

echo.
if %RESULT%==0 (
    echo ===============================================
    echo   UPLOAD SUCCESSFUL
    echo ===============================================
    echo The ATOM S3 will reboot into the new firmware.
    echo The COM port number may change after the reboot.
) else (
    echo ===============================================
    echo   UPLOAD FAILED (exit code %RESULT%)
    echo ===============================================
)
set EXIT_CODE=%RESULT%

:: ---- Post-flash: sync the freshly-generated BLE key ----
:: A full chip erase wipes LittleFS, so the device makes a NEW per-device
:: encryption key on first boot. Pull it over USB now so BLE recording /
:: variable sync work immediately without a manual profile upload. This
:: never changes the flasher's exit code — a sync miss isn't a flash
:: failure (the script prints its own guidance and the app re-syncs on the
:: next profile upload anyway).
if %RESULT%==0 (
    echo.
    echo -----------------------------------------------
    echo  Syncing BLE encryption key
    echo -----------------------------------------------
    python "%~dp0sync_key.py"
)

:end
echo.
pause
exit /b %EXIT_CODE%