#!/usr/bin/env python3
"""
Elemental Arena — one-command puck flasher.

Plug your puck in over USB (hold the BOOT button while connecting), then run:

    python3 flash.py

It installs the flash tool if needed, finds your device, writes the Elemental
Arena firmware, and tells you what to do next. No dev tools required.

Windows/Mac/Linux. If the firmware file isn't next to this script it is
downloaded automatically from the Arena.
"""
import glob
import os
import platform
import subprocess
import sys
import urllib.request

FW_NAME = "elemental-arena-2.2.24.bin"
BASE_URL = os.environ.get("ARENA_BASE_URL", "http://156.67.216.187:8004")
FW_URL = f"{BASE_URL}/static/flash/{FW_NAME}"
FLASH_OFFSET = "0x0"
CHIP = "esp32s3"
SETUP_URL = f"{BASE_URL}/setup"

C_OK = "\033[92m"; C_WARN = "\033[93m"; C_ERR = "\033[91m"; C_DIM = "\033[2m"; C_END = "\033[0m"
if platform.system() == "Windows":
    C_OK = C_WARN = C_ERR = C_DIM = C_END = ""


def say(msg, c=""):
    print(f"{c}{msg}{C_END}")


def ensure_esptool():
    try:
        import esptool  # noqa: F401
        return
    except ImportError:
        pass
    say("• Installing the flash tool (esptool)…", C_DIM)
    for args in (
        [sys.executable, "-m", "pip", "install", "--user", "--quiet", "esptool"],
        [sys.executable, "-m", "pip", "install", "--quiet", "esptool"],
    ):
        try:
            subprocess.check_call(args)
            import importlib
            importlib.invalidate_caches()
            import esptool  # noqa: F401
            return
        except Exception:
            continue
    say("Could not auto-install esptool. Run:  python3 -m pip install esptool", C_ERR)
    sys.exit(1)


def find_firmware():
    here = os.path.dirname(os.path.abspath(__file__))
    local = os.path.join(here, FW_NAME)
    if os.path.exists(local) and os.path.getsize(local) > 1_000_000:
        return local
    dest = os.path.join(here, FW_NAME)
    say(f"• Downloading firmware from {FW_URL} …", C_DIM)
    try:
        urllib.request.urlretrieve(FW_URL, dest)
    except Exception as e:
        say(f"Could not download firmware: {e}", C_ERR)
        say(f"Download it manually from {FW_URL} and put it next to this script.", C_WARN)
        sys.exit(1)
    return dest


def detect_port():
    # explicit override
    if len(sys.argv) > 1 and sys.argv[1].strip():
        return sys.argv[1].strip()
    env = os.environ.get("PUCK_PORT")
    if env:
        return env
    sysname = platform.system()
    patterns = []
    if sysname == "Darwin":
        patterns = ["/dev/cu.usbmodem*", "/dev/cu.wchusbserial*", "/dev/cu.SLAB_USBtoUART*", "/dev/cu.usbserial*"]
    elif sysname == "Linux":
        patterns = ["/dev/ttyACM*", "/dev/ttyUSB*"]
    hits = []
    for p in patterns:
        hits += glob.glob(p)
    hits = sorted(set(hits))
    if hits:
        return hits[0]
    return None  # let esptool auto-detect


def main():
    say("┌─────────────────────────────────────────────┐")
    say("│   ELEMENTAL ARENA — Puck Flasher            │")
    say("└─────────────────────────────────────────────┘")
    ensure_esptool()
    fw = find_firmware()
    say(f"• Firmware: {os.path.basename(fw)} ({os.path.getsize(fw)//1024} KB)", C_DIM)

    port = detect_port()
    if port:
        say(f"• Device port: {port}", C_DIM)
    else:
        say("• No port matched — letting the flasher auto-detect.", C_WARN)
        say("  If it can't connect: hold BOOT, tap RESET, release BOOT, and re-run.", C_WARN)

    cmd = [sys.executable, "-m", "esptool", "--chip", CHIP, "--baud", "921600"]
    if port:
        cmd += ["--port", port]
    cmd += ["--before", "default_reset", "--after", "hard_reset", "write_flash", FLASH_OFFSET, fw]

    say("\n• Flashing… (do not unplug — takes ~30–60s)\n", C_OK)
    try:
        rc = subprocess.call(cmd)
    except KeyboardInterrupt:
        say("\nCancelled.", C_WARN)
        sys.exit(1)

    if rc != 0:
        say("\n✗ Flashing failed.", C_ERR)
        say("  Fix: hold the BOOT button, tap RESET, release BOOT, then run this again.", C_WARN)
        say(f"  Still stuck? Pass the port explicitly:  python3 flash.py /dev/cu.usbmodemXXXX", C_WARN)
        sys.exit(rc)

    say("\n✓ Done! Your puck is now an Elemental Arena fighter.\n", C_OK)
    say("Next — follow along:", C_OK)
    say("  1. The puck reboots and shows a Wi-Fi setup screen (a hotspot + a code).")
    say("  2. On your phone/laptop, join that hotspot and pick your Wi-Fi.")
    say("  3. The puck then shows a 6-character PAIRING CODE.")
    say(f"  4. Open  {SETUP_URL}  , sign in, and enter that code to claim it.")
    say("  5. Choose your character — then say hi and challenge a friend!\n")
    say(f"{C_DIM}Arena: {BASE_URL}{C_END}")


if __name__ == "__main__":
    main()