"""Auto-detect a running ATOMS3 MacroPad and trigger ROM download mode.

Used by upload.bat to skip the manual "hold the side button while plugging
in" dance when a connected device is already running our firmware (which
exposes the `bootloader` serial command). On success, writes the detected
download-mode port to the file given as argv[1] and exits 0. On any
failure (no device, ping miss, bootloader command rejected, device fails
to re-enumerate) exits non-zero so the .bat falls back to manual mode.

This only works for devices already running our firmware — stock ATOMS3
firmware does not support the `bootloader` command, which is why
upload.bat keeps the manual flow as a fallback.
"""

import json
import sys
import time

import serial
import serial.tools.list_ports

DEVICE_ID = "ATOMS3-MACROPAD"
ESPRESSIF_VID = 0x303A
PING_TIMEOUT_S = 2.0
REENUMERATE_TIMEOUT_S = 8.0


def _espressif_ports() -> list:
    return [p for p in serial.tools.list_ports.comports() if p.vid == ESPRESSIF_VID]


def _ping(port: str) -> dict | None:
    """Open the port without resetting, send a ping, return parsed reply or None."""
    try:
        ser = serial.Serial()
        ser.port = port
        ser.baudrate = 115200
        ser.timeout = PING_TIMEOUT_S
        ser.dtr = False
        ser.rts = False
        ser.open()
    except (serial.SerialException, OSError):
        return None
    try:
        time.sleep(0.1)
        ser.dtr = True
        time.sleep(0.3)
        ser.reset_input_buffer()
        ser.write(b'{"cmd":"ping"}\n')
        line = ser.readline().decode("utf-8", errors="ignore").strip()
        if not line:
            return None
        return json.loads(line)
    except (serial.SerialException, json.JSONDecodeError, OSError):
        return None
    finally:
        try:
            ser.close()
        except Exception:
            pass


def _find_macropad() -> str | None:
    for p in _espressif_ports():
        reply = _ping(p.device)
        if reply and reply.get("id") == DEVICE_ID:
            print(f"  Found {DEVICE_ID} v{reply.get('ver', '?')} on {p.device}",
                  file=sys.stderr)
            return p.device
    return None


def _send_bootloader(port: str) -> bool:
    """Tell the firmware to call usb_persist_restart(RESTART_BOOTLOADER).

    The device acks with `{"rsp":"ok"}` then disappears off this port and
    re-enumerates as USB-Serial-JTAG (the native ESP32-S3 bootloader USB
    interface). We don't wait for the ack here beyond the read timeout —
    even a successful command vanishes the port within ~100 ms.
    """
    try:
        ser = serial.Serial()
        ser.port = port
        ser.baudrate = 115200
        ser.timeout = 1.0
        ser.dtr = False
        ser.rts = False
        ser.open()
    except (serial.SerialException, OSError) as exc:
        print(f"  ERROR: could not open {port}: {exc}", file=sys.stderr)
        return False
    try:
        time.sleep(0.1)
        ser.dtr = True
        time.sleep(0.3)
        ser.reset_input_buffer()
        ser.write(b'{"cmd":"bootloader"}\n')
        ser.flush()
        # Best-effort read of the ack; absence of one is fine — the chip
        # is already on its way to the ROM bootloader.
        try:
            line = ser.readline().decode("utf-8", errors="ignore").strip()
            print(f"  bootloader cmd reply: {line or '(none)'}", file=sys.stderr)
        except Exception:
            pass
        return True
    except (serial.SerialException, OSError) as exc:
        print(f"  ERROR sending bootloader command: {exc}", file=sys.stderr)
        return False
    finally:
        try:
            ser.close()
        except Exception:
            pass


def _wait_for_download_port(prev_macropad_port: str) -> str | None:
    """Wait for an Espressif device to appear in download mode.

    Strategy: poll the COM enumeration. Accept the first Espressif port we
    can see that either (a) is a *different* port than the one the firmware
    was on, or (b) is the same port but `ping` no longer answers (firmware
    is gone, USB-Serial-JTAG is up). The same-port case is common on
    Windows because the OS often reuses the COM number across the PHY
    switch.
    """
    deadline = time.monotonic() + REENUMERATE_TIMEOUT_S
    # Give the device a moment to vanish before we start polling.
    time.sleep(0.5)
    while time.monotonic() < deadline:
        ports = [p.device for p in _espressif_ports()]
        for port in ports:
            if port != prev_macropad_port:
                return port
        # Same port re-appeared (or never disappeared): confirm firmware is
        # no longer responding. If ping fails, this is the bootloader.
        if prev_macropad_port in ports:
            if _ping(prev_macropad_port) is None:
                return prev_macropad_port
        time.sleep(0.4)
    return None


def main(argv: list) -> int:
    if len(argv) < 2:
        print("usage: auto_enter_bootloader.py <output_port_file>", file=sys.stderr)
        return 2
    out_path = argv[1]

    print("Scanning for a running ATOMS3 MacroPad...", file=sys.stderr)
    fw_port = _find_macropad()
    if not fw_port:
        print("  No running MacroPad detected on any Espressif COM port.",
              file=sys.stderr)
        return 1

    print(f"Asking {fw_port} to enter download mode...", file=sys.stderr)
    if not _send_bootloader(fw_port):
        return 1

    print("Waiting for the device to re-enumerate as USB-Serial-JTAG...",
          file=sys.stderr)
    dl_port = _wait_for_download_port(fw_port)
    if not dl_port:
        print("  Timed out waiting for download-mode port.", file=sys.stderr)
        return 1

    print(f"  Download-mode port: {dl_port}", file=sys.stderr)
    try:
        with open(out_path, "w", encoding="ascii") as f:
            f.write(dl_port)
    except OSError as exc:
        print(f"  ERROR writing {out_path}: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))