#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""GD32 固件经 ESP32 中继刷写（tools/flash_gd32.py 删除后的等价实现）。

流程（与 ESP32 侧 http_handlers.cpp 的 ota_upload_handler /
ota_relay_trigger_handler 契约一致）：
  1. Release 的 Intel HEX 展平成从 0x08048000 起的连续 BIN（间隙填 0xFF）；
  2. POST /ota/upload，请求体 = {"total_size":N} 紧跟 BIN 字节；
  3. POST /ota/relay，请求体 = {"sha256":"<该 BIN 的 sha256>"}；设备侧还会做
     包身份防呆（复位向量必须落在 APP 区，误传 BOOT 会被拒）；
  4. 轮询 /system/status 直到 ota.sys_state=GD32_COMPLETE 且 ota.gd32_app 是
     期望版本、gd32_adc.uptime_sec 已归零；
  5. 把该 BIN 的 sha256 写进 lab/_last_gd32_sha.txt（板上镜像台账）。

用法: python scripts/flash_gd32.py [期望版本 如 1.6.5.0]
"""
import hashlib
import json
import os

import struct
import sys
import time
import urllib.error
import urllib.request

ESP = os.environ.get("DCU_ESP_IP", "192.168.31.127")
APP_BASE = 0x08048000
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", "..", ".."))
HEX = os.path.join(ROOT, "DCU_APP", "out", "GD32E508_DCU_APP", "GD32E508VE",
                   "Release", "GD32E508_DCU_APP.hex")
LEDGER = os.path.join(HERE, "..", "lab", "_last_gd32_sha.txt")


def hex_to_flat(path):
    """Intel HEX → 从 APP_BASE 起的连续字节。间隙填 0xFF（擦除态）。"""
    mem = {}
    base = 0
    with open(path, "r", encoding="ascii", errors="ignore") as f:
        for ln, line in enumerate(f, 1):
            line = line.strip()
            if not line.startswith(":"):
                continue
            raw = bytes.fromhex(line[1:])
            cnt, addr, rtype = raw[0], struct.unpack(">H", raw[1:3])[0], raw[3]
            data = raw[4:4 + cnt]
            if len(raw) != cnt + 5 or (sum(raw) & 0xFF) != 0:
                raise SystemExit(f"{path}:{ln} 校验和不对")
            if rtype == 0x00:
                for i, b in enumerate(data):
                    mem[base + addr + i] = b
            elif rtype == 0x02:
                base = struct.unpack(">H", data[:2])[0] << 4
            elif rtype == 0x04:
                base = struct.unpack(">H", data[:2])[0] << 16
            elif rtype == 0x01:
                break
    if not mem:
        raise SystemExit("HEX 里没有数据记录")
    lo, hi = min(mem), max(mem)
    if lo > APP_BASE:
        raise SystemExit(f"镜像起点 0x{lo:08X} 高于 APP 基址 0x{APP_BASE:08X}，不是 APP 包")
    if lo < APP_BASE:
        raise SystemExit(f"镜像含 0x{lo:08X} 的地址，落在 APP 区之外（疑似 BOOT 包）")
    out = bytearray(b"\xFF" * (hi - APP_BASE + 1))
    for a, b in mem.items():
        out[a - APP_BASE] = b
    return bytes(out)


def post(path, body, timeout=60):
    req = urllib.request.Request(f"http://{ESP}{path}", data=body, method="POST")
    req.add_header("Content-Type", "application/octet-stream")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, r.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode("utf-8", "replace")


def get(path, timeout=10):
    with urllib.request.urlopen(f"http://{ESP}{path}", timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8", "replace"))


def main():
    want = sys.argv[1] if len(sys.argv) > 1 else None
    if not os.path.isfile(HEX):
        raise SystemExit(f"找不到 {HEX}")
    blob = hex_to_flat(HEX)
    sha = hashlib.sha256(blob).hexdigest()
    print(f"镜像 {HEX}\n  展平 {len(blob)}B @0x{APP_BASE:08X}  sha256={sha}")

    hdr = json.dumps({"total_size": len(blob)}).encode("ascii")
    t0 = time.time()
    st, body = post("/ota/upload", hdr + blob, timeout=120)
    print(f"upload -> {st} {body[:200]}  ({time.time()-t0:.1f}s)")
    if st != 200:
        raise SystemExit("上传失败")

    t0 = time.time()
    st, body = post("/ota/relay", json.dumps({"sha256": sha}).encode("ascii"), 60)
    print(f"relay  -> {st} {body[:300]}  ({time.time()-t0:.1f}s)")

    deadline = time.time() + 90
    last = ""
    while time.time() < deadline:
        time.sleep(2)
        try:
            s = get("/system/status")
        except Exception as e:
            print(f"  等待中… status 读失败 {type(e).__name__}")
            continue
        o = s.get("ota", {})
        line = (f"  sys_state={o.get('sys_state')} gd32_app={o.get('gd32_app')} "
                f"alive={o.get('gd32_app_alive')} uptime={s.get('uptime_sec')}")
        if line != last:
            print(line)
            last = line
        if o.get("sys_state") == "GD32_COMPLETE" and o.get("gd32_app"):
            ver = str(o["gd32_app"]).lstrip("v")
            if want and ver != want:
                raise SystemExit(f"板上版本 {o['gd32_app']} 与期望 {want} 不符 —— 刷写未生效")
            with open(LEDGER, "w", encoding="ascii") as f:
                f.write(sha)
            print(f"完成: 板上 gd32_app={o['gd32_app']}，台账已写 {os.path.normpath(LEDGER)}")
            return
    raise SystemExit("90s 内未见 GD32_COMPLETE —— 刷写未完成，别继续测试")


if __name__ == "__main__":
    main()