#!/usr/bin/env python3
"""SUOTA 全量烧录执行器（用户授权 2026-07-29，范围见 DECISIONS.md）。
全程事件流写 docs/suota-flash-log.jsonl；结束打印统计与事后 DIS 验证。
用法：.venv/bin/python tools/flash_run.py [image] [gpio_hex]
"""
import asyncio
import datetime
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, "src")
from ppclock.bridge import BridgeTransport  # noqa: E402

HOST, PORT = "192.168.61.35", 8971
DEVICE = "18:BC:5A:5D:BF:28"
IMAGE = sys.argv[1] if len(sys.argv) > 1 else "analysis/firmware/fw_tst.img"
GPIO = sys.argv[2] if len(sys.argv) > 2 else "05060040"  # BE 序 C（gpio_probe 实证）
LOG = Path("docs/suota-flash-log.jsonl")


def log(entry):
    entry["ts"] = datetime.datetime.now().isoformat(timespec="milliseconds")
    with LOG.open("a", encoding="utf-8") as f:
        f.write(json.dumps(entry, ensure_ascii=False) + "\n")


async def ota_with_gpio(host, port, image, gpio_hex, prog):
    """带 gpio_map 参数的 run_ota（BridgeTransport.run_ota 固定默认，临时直连 RPC）。"""
    reader, writer = await asyncio.open_connection(host, port)
    events = []

    async def read_loop():
        async for line in reader:
            msg = json.loads(line)
            events.append(msg)
            prog(msg)

    task = asyncio.create_task(read_loop())
    try:
        writer.write(json.dumps({"id": 1, "op": "ota", "image": image.hex(),
                                 "gpio_map": gpio_hex}).encode() + b"\n")
        await writer.drain()
        while True:
            done, _ = await asyncio.wait({task}, timeout=1800)
            if done:
                break
            # 检查是否已收到最终响应
            finals = [e for e in events if "id" in e]
            if finals:
                break
            if task.done():
                break
        finals = [e for e in events if "id" in e]
        return finals[-1] if finals else {"ok": False, "error": "no final response"}
    finally:
        task.cancel()
        writer.close()


async def main():
    LOG.unlink(missing_ok=True)
    log({"kind": "run_start", "image": IMAGE, "gpio": GPIO,
         "authorization": "用户授权 2026-07-29（DECISIONS.md）"})
    image = Path(IMAGE).read_bytes()
    t0 = time.time()
    stats = {}

    def prog(msg):
        el = time.time() - t0
        entry = {"el": round(el, 2), "msg": msg}
        log({"kind": "event", **entry})
        e = msg.get("event")
        if e == "ota_stage":
            for c in msg.get("codes", []):
                stats[c] = stats.get(c, 0) + 1
            print(f"[{el:6.1f}s] 阶段: {msg['stage']}"
                  + (f" codes={msg['codes']}" if msg.get("codes") else ""), flush=True)
        elif e == "ota_progress":
            stats[2] = stats.get(2, 0) + 1
            if msg["block"] % 20 == 0 or msg["block"] == msg["total"]:
                print(f"[{el:6.1f}s] 块确认(2) {msg['block']}/{msg['total']}", flush=True)

    # 守候清醒窗口（设备无唤醒按钮、自动周期醒）
    for attempt in range(20):
        try:
            async with BridgeTransport(address=DEVICE, timeout=30, subscribe=False,
                                       host=HOST, port=PORT) as t:
                print(f"已连接（第 {attempt + 1} 次尝试），开始全量烧录...", flush=True)
                log({"kind": "connected", "device": DEVICE, "attempt": attempt + 1})
                result = await ota_with_gpio(HOST, PORT, image, GPIO, prog)
                print(f"[{time.time() - t0:6.1f}s] 最终结果: {result}", flush=True)
                log({"kind": "run_end", "result": result, "stats": stats})
                break
        except Exception as e:
            print(f"  等待设备唤醒...({attempt + 1}/20) {type(e).__name__}", flush=True)
            await asyncio.sleep(15)
    else:
        print("设备持续未醒，放弃", flush=True)
        log({"kind": "run_abort", "reason": "device_never_awake"})
        return
    print(f"=== 状态统计: {stats} ===", flush=True)


asyncio.run(main())