# -*- coding: utf-8 -*-
"""Flash Boot then APP via cspybat (AT32F403AC), reset, verify Boot strings."""
from __future__ import annotations

import os
import shutil
import subprocess
import time
from pathlib import Path

BASE = Path(r"C:\Users\qjyu\Documents\SoundWalker\涓€璇哄浗闄呭悏浠朶Code\YNGJ-GT1-M - AT32F403ARCT7")
STAGE = Path(r"C:\Temp\k1flash_boot")
CSPY = Path(r"C:\Program Files (x86)\IAR Systems\Embedded Workbench 7.3\common\bin\cspybat.exe")
JLINK = Path(r"C:\Program Files\SEGGER\JLink_V818\JLink.exe")

BOOT_OUT = BASE / "AT32F403ARCT7_BOOT" / "project" / "IAR_V7.4" / "AT32F403ARCT7_BOOT" / "Exe" / "AT32F403ARCT7_BOOT.out"
APP_OUT = BASE / "project" / "IAR_V7.4" / "YNGJ-GT1-M" / "Exe" / "YNGJ-GT1-M.out"
GEN_TMPL = BASE / "project" / "IAR_V7.4" / "settings" / "YNGJ-GT1-M.YNGJ-GT1-M.general.xcl"
DRV_TMPL = BASE / "project" / "IAR_V7.4" / "settings" / "YNGJ-GT1-M.YNGJ-GT1-M.driver.xcl"
RESET_JLINK = BASE / "tools" / "reset_run.jlink"

NEW_GBK = bytes([0xC9, 0xD5, 0xC2, 0xBC, 0xC4, 0xA3, 0xCA, 0xBD])  # 鐑у綍妯″紡
OLD_GBK = bytes([0xC9, 0xFD, 0xBC, 0xB6, 0xC4, 0xA3, 0xCA, 0xBD])  # 鍗囩骇妯″紡


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


def stage_xcl(out_file: Path) -> tuple[Path, Path]:
    STAGE.mkdir(parents=True, exist_ok=True)
    staged_out = STAGE / out_file.name
    shutil.copy2(out_file, staged_out)

    gen_lines = GEN_TMPL.read_text(encoding="utf-8", errors="replace").splitlines()
    new_gen = []
    for ln in gen_lines:
        if ".out" in ln and ("YNGJ" in ln or "BOOT" in ln or "AT32" in ln):
            new_gen.append(f'"{staged_out}" ')
        else:
            new_gen.append(ln)
    gen_path = STAGE / "general.xcl"
    gen_path.write_text("\n".join(new_gen) + "\n", encoding="utf-8", newline="\n")

    drv_lines = DRV_TMPL.read_text(encoding="utf-8", errors="replace").splitlines()
    forced = [ln for ln in drv_lines if not ln.strip().startswith("--jlink_device")]
    forced.append("--jlink_device=AT32F403AC")
    drv_path = STAGE / "driver.xcl"
    drv_path.write_text("\n".join(forced) + "\n", encoding="utf-8", newline="\n")
    return gen_path, drv_path, staged_out


def cspy_download(out_file: Path, tag: str) -> None:
    if not out_file.is_file():
        raise SystemExit(f"missing {out_file}")
    gen_path, drv_path, staged_out = stage_xcl(out_file)
    logp = STAGE / f"cspy_{tag}.log"
    cmd = [
        str(CSPY),
        "-f",
        str(gen_path),
        f"--debug_file={staged_out}",
        "--download_only",
        "--backend",
        "-f",
        str(drv_path),
    ]
    print(f"cspybat download {tag}: {staged_out.name}", flush=True)
    with logp.open("w", encoding="utf-8", errors="replace") as log:
        r = subprocess.run(cmd, stdout=log, stderr=subprocess.STDOUT, timeout=240)
    text = logp.read_text(encoding="utf-8", errors="replace")
    print(text[-1500:], flush=True)
    if r.returncode != 0:
        raise SystemExit(f"cspybat {tag} failed rc={r.returncode}")


def jlink_reset() -> None:
    print("J-Link reset/run...", flush=True)
    subprocess.run(
        [
            str(JLINK),
            "-Device",
            "AT32F403AC",
            "-If",
            "SWD",
            "-Speed",
            "4000",
            "-AutoConnect",
            "1",
            "-CommandFile",
            str(RESET_JLINK),
        ],
        capture_output=True,
        timeout=45,
    )
    time.sleep(2.5)


def verify_boot_in_mcu() -> None:
    import pylink

    j = pylink.JLink()
    j.open()
    try:
        try:
            j.exec_command("HideDeviceSelection = 1")
        except Exception:
            pass
        j.set_tif(pylink.enums.JLinkInterfaces.SWD)
        last = None
        for dev in ("AT32F403AC", "Cortex-M4", "AT32F403A"):
            try:
                try:
                    j.exec_command(f"Device = {dev}")
                except Exception:
                    pass
                j.connect(dev)
                print(f"verify connect: {dev}", flush=True)
                break
            except Exception as exc:
                last = exc
        else:
            raise SystemExit(f"verify connect failed: {last}")

        j.halt()
        # Boot image is in first ~32KB; search GBK markers
        chunk = bytes(j.memory_read8(0x08000000, 0x8000))
        has_new = NEW_GBK in chunk
        has_old = OLD_GBK in chunk
        print(f"MCU@0x08000000 contains 鐑у綍妯″紡={has_new} 鍗囩骇妯″紡={has_old}", flush=True)
        if not has_new:
            raise SystemExit("Boot flash verify FAILED: new 鐑у綍妯″紡 string missing")
        if has_old:
            print("WARN: old 鍗囩骇妯″紡 string still present somewhere in Boot region", flush=True)
        else:
            print("Boot flash verify OK", flush=True)
        j.reset(halt=False)
    finally:
        j.close()


def main() -> None:
    boot_bin = BOOT_OUT.with_suffix(".bin")
    data = boot_bin.read_bytes()
    print(
        f"local BOOT.bin: new={NEW_GBK in data} old={OLD_GBK in data} size={len(data)}",
        flush=True,
    )
    kill_debuggers()
    cspy_download(BOOT_OUT, "boot")
    kill_debuggers()
    cspy_download(APP_OUT, "app")
    jlink_reset()
    verify_boot_in_mcu()
    print("Done. Keep KEY for burn mode; screen should use UI0902_FLASH_MODE if ExtFlash has art.", flush=True)


if __name__ == "__main__":
    main()