# -*- coding: utf-8 -*-
"""Flash K1 MCU APP (J-Link) + ExtFlash ALL.res (RTT flash all).

Usage:
  python tools/flash_mcu_and_extflash.py
  python tools/flash_mcu_and_extflash.py --mcu PATH.bin --all PATH.res
"""
from __future__ import annotations

import argparse
import shutil
import subprocess
import time
from pathlib import Path

PROJ = Path(__file__).resolve().parents[1]
REPO = PROJ.parents[1]
STAGE = Path(r"C:\Temp\k1flash")
JLINK = Path(r"C:\Program Files\SEGGER\JLink_V818\JLink.exe")
DEFAULT_DEVICES = ("AT32F403AC", "Cortex-M4", "AT32F403A")

DEFAULT_MCU = (
    REPO
    / "tools"
    / "out"
    / "K1_MCU_v0.2.11_toneRes_v0914_20260914_d06e44b"
    / "YNGJ-GT1-M_MCU_v0.2.11_20260914.bin"
)
FALLBACK_MCU = PROJ / "project" / "IAR_V7.4" / "YNGJ-GT1-M" / "Exe" / "YNGJ-GT1-M.bin"

DEFAULT_ALL = REPO / "tools" / "out" / "extflash_ALL_tone0914_ui0902.res"
FALLBACK_ALL = (
    REPO
    / "tools"
    / "out"
    / "K1_MCU_v0.2.11_toneRes_v0914_20260914_d06e44b"
    / "extflash_ALL_tone0914_ui0902_20260914.res"
)


def kill_debuggers() -> None:
    for n in (
        "cspybat.exe",
        "CSpyBat.exe",
        "JLink.exe",
        "JLinkGUIServer.exe",
        "JLinkRTTClient.exe",
        "JFlash.exe",
    ):
        subprocess.run(["taskkill", "/F", "/IM", n], capture_output=True)
    time.sleep(0.8)


def resolve_path(primary: Path, fallback: Path) -> Path:
    if primary.is_file():
        return primary
    if fallback.is_file():
        return fallback
    raise SystemExit(f"missing file:\n  {primary}\n  {fallback}")


def flash_mcu_jlink(mcu_bin: Path) -> None:
    STAGE.mkdir(parents=True, exist_ok=True)
    staged = STAGE / "YNGJ-GT1-M.bin"
    shutil.copy2(mcu_bin, staged)
    cmdfile = STAGE / "flash_app.jlink"
    cmdfile.write_text(
        "\n".join(
            [
                "HideDeviceSelection 1",
                "si SWD",
                "speed 4000",
                "device AT32F403AC",
                "connect",
                "halt",
                f"loadbin {staged}, 0x08008000",
                "r",
                "g",
                "exit",
                "",
            ]
        ),
        encoding="ascii",
    )
    if not JLINK.is_file():
        raise SystemExit(f"missing JLink.exe: {JLINK}")
    print(f"[1/3] J-Link flash MCU → 0x08008000 ({staged.name}, {staged.stat().st_size} bytes)")
    r = subprocess.run(
        [
            str(JLINK),
            "-Device",
            "AT32F403AC",
            "-If",
            "SWD",
            "-Speed",
            "4000",
            "-AutoConnect",
            "1",
            "-CommandFile",
            str(cmdfile),
        ],
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=120,
    )
    out = (r.stdout or "") + (r.stderr or "")
    print(out[-1200:] if out else f"rc={r.returncode}")
    if r.returncode != 0:
        raise SystemExit(f"J-Link MCU flash failed rc={r.returncode}")
    if "Failed" in out and "O.K." not in out:
        raise SystemExit("J-Link MCU flash reported failure")
    time.sleep(2.5)


def connect_jlink(device: str):
    import pylink

    jlink = pylink.JLink()
    jlink.open()
    try:
        jlink.exec_command("HideDeviceSelection = 1")
    except Exception:
        pass
    jlink.set_tif(pylink.enums.JLinkInterfaces.SWD)
    last = None
    for candidate in (device, *DEFAULT_DEVICES):
        try:
            try:
                jlink.exec_command(f"Device = {candidate}")
            except Exception:
                pass
            jlink.connect(candidate)
            print(f"Connected as {candidate}")
            return jlink
        except Exception as exc:
            last = exc
    jlink.close()
    raise SystemExit(f"J-Link connect failed: {last}")


def find_rtt_cb(jlink, ram_base: int = 0x20000000, ram_size: int = 0x18000):
    needle = b"SEGGER RTT"
    chunk = 0x1000
    was_halted = jlink.halted()
    if not was_halted:
        jlink.halt()
    try:
        for off in range(0, ram_size, chunk):
            data = bytes(jlink.memory_read8(ram_base + off, min(chunk, ram_size - off)))
            idx = data.find(needle)
            if idx >= 0:
                return ram_base + off + idx
    finally:
        if not was_halted:
            if hasattr(jlink, "restart"):
                jlink.restart()
            else:
                jlink.go()
            time.sleep(0.1)
    return None


def rtt_read_text(jlink, timeout_s: float = 0.2) -> str:
    deadline = time.time() + timeout_s
    chunks: list[bytes] = []
    while time.time() < deadline:
        try:
            data = bytes(jlink.rtt_read(0, 512) or [])
        except Exception:
            data = b""
        if data:
            chunks.append(data)
            deadline = time.time() + timeout_s
        else:
            time.sleep(0.02)
    return b"".join(chunks).decode("ascii", errors="replace")


def wait_for(jlink, marker: str, timeout_s: float) -> str:
    deadline = time.time() + timeout_s
    buf = ""
    while time.time() < deadline:
        buf += rtt_read_text(jlink, 0.15)
        if marker in buf:
            return buf
    raise SystemExit(f"Timeout waiting for {marker!r}\n--- RTT ---\n{buf[-1000:]}")


def send_bytes(jlink, payload: bytes) -> None:
    off = 0
    while off < len(payload):
        try:
            n = jlink.rtt_write(0, list(payload[off : off + 64]))
        except Exception:
            time.sleep(0.05)
            continue
        if n is None or n <= 0:
            time.sleep(0.02)
            continue
        off += n


def flash_all_rtt(all_res: Path, device: str) -> None:
    data = all_res.read_bytes()
    size = len(data)
    print(f"[2/3] RTT flash ALL.res @ 0x0 ({all_res.name}, {size} bytes)")

    jlink = connect_jlink(device)
    try:
        print("hardware reset...")
        jlink.reset(halt=False)
        time.sleep(4.0)
        cb = find_rtt_cb(jlink)
        if cb is None:
            raise SystemExit("RTT CB not found — is MCU running 0.2.11+ with APP_LOG?")
        print(f"RTT CB @ 0x{cb:08X}")
        jlink.rtt_start(cb)
        time.sleep(0.5)
        _ = rtt_read_text(jlink, 0.5)

        send_bytes(jlink, f"flash all {size}\n".encode("ascii"))
        log = wait_for(jlink, "FLASH_ALL_GO", 20.0)
        print(log.strip().splitlines()[-1])
        time.sleep(0.2)

        chunk = 256
        sent = 0
        t0 = time.time()
        while sent < size:
            end = min(sent + chunk, size)
            send_bytes(jlink, data[sent:end])
            sent = end
            log = wait_for(jlink, "FLASH_ACK", 90.0)
            if "FLASH_ALL_OK" in log:
                print(log.strip().splitlines()[-1])
                print(f"ALL.res done in {time.time() - t0:.1f}s")
                break
            if sent % (128 * 1024) == 0 or sent == size:
                print(f"  {sent}/{size} ({100.0 * sent / size:.1f}%) {time.time() - t0:.1f}s")
        else:
            log = wait_for(jlink, "FLASH_ALL_OK", 90.0)
            print(log.strip().splitlines()[-1])

        print("[3/3] verify BIN3 via tone status")
        send_bytes(jlink, b"tone status\n")
        log = wait_for(jlink, "TONE @BIN3", 10.0)
        for line in log.strip().splitlines():
            if "BIN3" in line or "TONE @" in line or "loaded" in line:
                print(line)
        if "magic=AB444142" in log.replace(" ", "") or "magic=AB" in log:
            # formats like magic=AB444142 or magic=AB44...
            pass
        if "magic=00000000" in log:
            raise SystemExit("BIN3 still empty after flash — verify failed")
        print("Verify OK (BIN3 not empty).")
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mcu", type=Path, default=DEFAULT_MCU)
    ap.add_argument("--all", dest="all_res", type=Path, default=DEFAULT_ALL)
    ap.add_argument("--device", default="AT32F403AC")
    ap.add_argument("--skip-mcu", action="store_true")
    args = ap.parse_args()

    mcu = resolve_path(args.mcu, FALLBACK_MCU)
    all_res = resolve_path(args.all_res, FALLBACK_ALL)
    print(f"MCU : {mcu}")
    print(f"ALL : {all_res}")

    kill_debuggers()
    if not args.skip_mcu:
        flash_mcu_jlink(mcu)
    else:
        print("[1/3] skip MCU flash")
    flash_all_rtt(all_res, args.device)
    print("Done. Please power-cycle / long-press power and test universal pick.")


if __name__ == "__main__":
    main()