#!/usr/bin/env python3
"""Flash this firmware onto another ESP32-P4 board over USB, then give it its own identity.

Steps: pick the port, check the chip (ESP32-P4, revision the image supports, flash size),
optionally erase, write bootloader + partition table + app, wait for the board's IP in its boot
log, then set hostname / stream name / multicast address over the web API so it doesn't clash
with other boards on the network.

Needs esptool (in the ESP-IDF Python env, or `pip install esptool`) and a build (`idf.py build`).
Every question can also be answered with an option, e.g.:
  tools/flash_board.py --port /dev/ttyACM0 --erase --hostname p4-aes67-2 --name "P4 AES67 2" \\
                       --mcast 239.69.1.11
Boards already running this firmware and on the network update over the web (OTA) instead.
"""
import argparse
import json
import os
import re
import struct
import subprocess
import sys
import time
import urllib.request

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MIN_FLASH_MB = 32   # the bootloader header says 32 MB (sdkconfig); smaller chips aren't supported


def die(msg):
    print(f"\nERROR: {msg}")
    sys.exit(1)


def ask(prompt, default=None):
    s = input(f"{prompt}{f' [{default}]' if default is not None else ''}: ").strip()
    return s or (default if default is not None else "")


def yes(prompt, default):
    s = input(f"{prompt} [{'Y/n' if default else 'y/N'}]: ").strip().lower()
    return default if not s else s.startswith("y")


def esptool(port, *args, capture=True):
    cmd = [sys.executable, "-m", "esptool", "--chip", "esp32p4", "--port", port, *args]
    if not capture:
        if subprocess.call(cmd) != 0:
            die("esptool failed (see above)")
        return ""
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        die(f"esptool {' '.join(args)} failed:\n{r.stdout}{r.stderr}")
    return r.stdout


def pick_port(given):
    if given:
        return given
    try:
        from serial.tools import list_ports
    except ImportError:
        die("pyserial missing: run inside the ESP-IDF environment or `pip install esptool`")
    ports = [p for p in list_ports.comports() if "ACM" in p.device or "USB" in p.device]
    if not ports:
        die("no USB serial port found: connect the board's USB-C port (not the PoE/Ethernet side)")
    for i, p in enumerate(ports, 1):
        print(f"  {i}) {p.device}  {p.description}")
    if len(ports) == 1:
        return ports[0].device if yes(f"Use {ports[0].device}?", True) else die("no port chosen")
    n = ask("Port number", "1")
    return ports[int(n) - 1].device


def load_images(build):
    """(offset, path) pairs and flash options from build/flash_args."""
    fa = os.path.join(build, "flash_args")
    if not os.path.exists(fa):
        die(f"{fa} not found: build first (idf.py build)")
    lines = open(fa).read().split("\n")
    opts = lines[0].split()
    images = []
    for line in lines[1:]:
        if line.strip():
            off, path = line.split()
            images.append((off, os.path.join(build, path)))
    return opts, images


def image_rev_range(app):
    """Chip revision range from the app image header (esp_image_header_t), as (min, max) * 100."""
    h = open(app, "rb").read(24)
    if h[0] != 0xE9:
        die(f"{app} is not an ESP app image")
    min_full, max_full = struct.unpack_from("<HH", h, 15)
    return min_full, max_full


def app_version(app):
    # esp_app_desc_t after the image header (24 B) and first segment header (8 B): version at +16.
    d = open(app, "rb").read(24 + 8 + 256)[32:]
    return d[16:48].split(b"\0")[0].decode(errors="replace")


def check_chip(port, app):
    out = esptool(port, "flash_id")
    m = re.search(r"ESP32-P4 \(revision v(\d+)\.(\d+)\)", out)
    if not m:
        die(f"not an ESP32-P4?\n{out}")
    rev = int(m.group(1)) * 100 + int(m.group(2))
    lo, hi = image_rev_range(app)
    print(f"  chip: ESP32-P4 revision v{rev // 100}.{rev % 100}")
    if not lo <= rev <= hi:
        die(f"this build supports chip revisions v{lo // 100}.{lo % 100} to v{hi // 100}.{hi % 100}; "
            f"the board is v{rev // 100}.{rev % 100}. Build with the matching "
            "CONFIG_ESP32P4_SELECTS_REV_LESS_V3 setting (see CLAUDE.md).")
    m = re.search(r"Detected flash size: (\d+)MB", out)
    if not m:
        die(f"could not read the flash size:\n{out}")
    mb = int(m.group(1))
    print(f"  flash: {mb} MB")
    if mb < MIN_FLASH_MB:
        die(f"{mb} MB flash; this build needs {MIN_FLASH_MB} MB")


def wait_for_ip(port, timeout=90):
    """Read the boot log until the board reports its IP (the port open may reset it: fine here)."""
    import serial
    print(f"  waiting for the board's IP (boot log on {port}, up to {timeout} s; Ethernet plugged in?)")
    end = time.time() + timeout
    with serial.Serial(port, 115200, timeout=1) as s:
        buf = b""
        while time.time() < end:
            buf += s.read(4096)
            m = re.search(rb"got IP (\d+\.\d+\.\d+\.\d+)", buf)
            if m:
                return m.group(1).decode()
            buf = buf[-4096:]
    return None


def api(ip, method, path, body=None):
    req = urllib.request.Request(f"http://{ip}{path}", method=method,
                                 data=json.dumps(body).encode() if body is not None else None,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=10) as r:
        t = r.read()
        return json.loads(t) if t else {}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--port")
    ap.add_argument("--build", default=os.path.join(REPO, "build"), help="build directory (default: build/)")
    ap.add_argument("--erase", action=argparse.BooleanOptionalAction,
                    help="erase the whole flash first (clears any stored config)")
    ap.add_argument("--hostname")
    ap.add_argument("--name", help="stream / Spotify device name")
    ap.add_argument("--mcast", help="AES67 multicast address")
    ap.add_argument("--no-config", action="store_true", help="flash only, keep the default identity")
    a = ap.parse_args()

    opts, images = load_images(a.build)
    app = next(p for off, p in images if p.endswith("aes67_p4.bin"))
    print(f"Firmware: {app_version(app)}  ({app})\n")

    print("1) Port")
    port = pick_port(a.port)

    print("\n2) Checking the board")
    check_chip(port, app)

    print("\n3) Flashing")
    erase = a.erase if a.erase is not None else yes(
        "  Erase the whole flash first? (yes for a new board or to clear an old config)", True)
    if erase:
        esptool(port, "erase_flash", capture=False)
    args = ["write_flash", *opts]
    for off, path in images:
        args += [off, path]
    esptool(port, *args, capture=False)

    if a.no_config:
        print("\nDone (default identity: hostname p4-aes67, multicast 239.69.1.10).")
        return

    print("\n4) Finding the board on the network")
    ip = wait_for_ip(port)
    if not ip:
        die("no IP seen in the boot log. Check the Ethernet cable/DHCP; the board is flashed, set the "
            "identity later in the web UI.")
    print(f"  board is at {ip}")
    for _ in range(30):     # web server comes up right after the IP
        try:
            cfg = api(ip, "GET", "/api/config")
            break
        except OSError:
            time.sleep(1)
    else:
        die(f"web UI at http://{ip}/ does not answer")

    print("\n5) Identity (must differ from every other board on this network)")
    host = a.hostname or ask("  Hostname", "p4-aes67-2")
    name = a.name or ask("  Stream / Spotify device name", "P4 AES67 2")
    mcast = a.mcast or ask("  AES67 multicast address", "239.69.1.11")
    cfg["net"]["hostname"] = host
    cfg["aes67"]["name"] = name
    cfg["aes67"]["mcast"] = mcast
    cfg["source"]["spotify_name"] = name
    try:
        api(ip, "POST", "/api/config", cfg)
    except urllib.error.HTTPError as e:
        die(f"config rejected: {e.read().decode()}")
    print(f"\nDone. Web UI: http://{ip}/  or  http://{host}.local/")
    print("Still to set in the web UI if needed: Spotify client ID/secret, source mode, PTP role.")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\naborted")
        sys.exit(1)