#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""音师自助：把投放目录的 1/2/3.bin 打成 extflash_ALL_*.res，供 SoundWalkerIAP USB 烧录。

不走 J-Link；不改固件头文件；不重编 MCU。
需要本机已有 ui0902_res.bin（或最近一次发布包里的 ALL.res 可拆出 UI 段）。
"""
from __future__ import annotations

import argparse
import shutil
import struct
import sys
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]  # firmware project
REPO = ROOT.parent.parent  # 一诺国际吉他
OUT_DIR = ROOT / "tools" / "out"
DROP_DIR = REPO / "Doc" / "音色文件" / "音师投放"
IAP_DIR = REPO / "升级" / "MCU主控升级"
UI0902_BIN = OUT_DIR / "ui0902_res.bin"
REPO_OUT = REPO / "tools" / "out"

OFF_LOGO = 0x00000000
OFF_CHARGING = 0x0000CB70
OFF_BIN1 = 0x0001B8F0
OFF_BIN2 = 0x0009D07D
OFF_BIN3 = 0x000A71AC
UI0902_RES_BASE = 0x00100000
DAB_MAGIC = b"\xABDAB"
MAX_SONG_BIN_BYTES = 41 * 1024
BIN2_SLOT_BYTES = OFF_BIN3 - OFF_BIN2


def find_ziliao() -> Path:
    for p in REPO.iterdir():
        if p.is_dir() and (p / "logo.bin").is_file() and (p / "Charg.bin").is_file():
            return p
    raise SystemExit("找不到 资料/logo.bin 与 Charg.bin，无法打包。")


def dab_ok(blob: bytes, off: int, expect_cnt: int, label: str) -> None:
    if off + 16 > len(blob):
        raise SystemExit(f"{label} @0x{off:X}: 文件太短，无法校验")
    magic = blob[off : off + 4]
    cnt = struct.unpack_from("<I", blob, off + 12)[0]
    if magic != DAB_MAGIC:
        raise SystemExit(
            f"{label} @0x{off:X}: DAB 魔数错误 {magic.hex()}（期望 {DAB_MAGIC.hex()}），请确认文件是否正确导出"
        )
    if cnt != expect_cnt:
        raise SystemExit(
            f"{label} @0x{off:X}: 预设数量={cnt}，期望 {expect_cnt}。"
            f"{'若曲目数量有变，需找开发同步固件曲名表。' if label.startswith('2.bin') else ''}"
        )


def require_equals(blob: bytes, off: int, src: bytes, label: str) -> None:
    end = off + len(src)
    if end > len(blob) or blob[off:end] != src:
        raise SystemExit(f"内部校验失败：{label} 未正确写入资源包")


def find_ui_baseline() -> bytes:
    if UI0902_BIN.is_file():
        data = UI0902_BIN.read_bytes()
        print(f"UI 底图: {UI0902_BIN} ({len(data)} bytes)")
        return data

    # Fallback: newest release ALL.res under repo tools/out
    candidates: list[Path] = []
    if REPO_OUT.is_dir():
        candidates.extend(REPO_OUT.glob("**/extflash_ALL_*.res"))
        candidates.extend(REPO_OUT.glob("extflash_ALL_*.res"))
    candidates = sorted({p.resolve() for p in candidates if p.is_file()}, key=lambda p: p.stat().st_mtime, reverse=True)
    for all_path in candidates:
        raw = all_path.read_bytes()
        if len(raw) > UI0902_RES_BASE:
            ui = raw[UI0902_RES_BASE:]
            print(f"UI 底图: 从发布包拆出 {all_path.name} @0x{UI0902_RES_BASE:X} ({len(ui)} bytes)")
            return ui

    raise SystemExit(
        "缺少 UI 底图 ui0902_res.bin。\n"
        "请向开发索取一次含 UI 的基线（tools/out/ui0902_res.bin，或完整 extflash_ALL_*.res），"
        "放到工程 tools/out/ 后再运行本工具。"
    )


def pack_tone(bin1: bytes, bin2: bytes, bin3: bytes, logo: bytes, charg: bytes) -> bytes:
    if len(bin2) > MAX_SONG_BIN_BYTES:
        raise SystemExit(f"2.bin 过大: {len(bin2)} > {MAX_SONG_BIN_BYTES} (41KB 上限)")
    if len(bin2) > BIN2_SLOT_BYTES:
        raise SystemExit(
            f"2.bin 过大: {len(bin2)} > 槽位 {BIN2_SLOT_BYTES}，会挤占 3.bin 固定地址 0x{OFF_BIN3:X}"
        )

    parts = [
        (OFF_LOGO, logo, "logo"),
        (OFF_CHARGING, charg, "Charg"),
        (OFF_BIN1, bin1, "1.bin"),
        (OFF_BIN2, bin2, "2.bin"),
        (OFF_BIN3, bin3, "3.bin"),
    ]
    blobs: list[bytes] = []
    cursor = 0
    for force_off, data, name in parts:
        if cursor > force_off:
            raise SystemExit(f"{name}: 前一段过大，无法放到 0x{force_off:X}")
        if cursor < force_off:
            blobs.append(b"\xFF" * (force_off - cursor))
            cursor = force_off
        blobs.append(data)
        cursor += len(data)

    packed = b"".join(blobs)
    dab_ok(packed, OFF_BIN1, 31, "1.bin(节奏)")
    dab_ok(packed, OFF_BIN2, 1, "2.bin(本地曲目)")
    dab_ok(packed, OFF_BIN3, 3, "3.bin(万能)")
    require_equals(packed, OFF_BIN1, bin1, "1.bin")
    require_equals(packed, OFF_BIN2, bin2, "2.bin")
    require_equals(packed, OFF_BIN3, bin3, "3.bin")
    if cursor > UI0902_RES_BASE:
        raise SystemExit("音色区溢出到 UI 区 (0x100000)")
    return packed


def write_flash_guide(path: Path, all_name: str, stamp: str) -> None:
    text = "\n".join(
        [
            "音师音色自助包 — 请这样用 SoundWalkerIAP 刷机",
            f"生成时间: {stamp}",
            "",
            "==== 本次只需刷外部 Flash（不用重刷 Boot/APP，除非开发另有说明） ====",
            f"  文件: {all_name}",
            "  地址: 0x00000000（从外部 Flash 开头整包写入）",
            "",
            "步骤:",
            "  1. 吉他进入 USB 升级模式（参见 升级步骤.docx）",
            "  2. 打开本目录的 SoundWalkerIAP.exe",
            "  3. 先整片擦除外部 Flash",
            f"  4. 选择 {all_name}，烧录到外部 Flash @ 0x00000000",
            "  5. 退出升级模式，重启，试听节奏/本地曲/万能",
            "",
            "注意:",
            "  - 必须刷本包 ALL.res，不要只刷单独的 1/2/3.bin",
            "  - 只刷音色、不带 UI 会导致模式选择页花屏",
            "  - 2.bin 必须 ≤41KB；3.bin 地址固定，不可随 2.bin 变长漂移",
            "  - 若增删本地曲目或改曲名，需找开发同步固件后再测",
            "",
        ]
    )
    path.write_text(text, encoding="utf-8", newline="\n")


def main() -> int:
    ap = argparse.ArgumentParser(description="音师自助打包 1/2/3.bin → ALL.res（USB IAP）")
    ap.add_argument(
        "--drop",
        type=Path,
        default=DROP_DIR,
        help=f"投放目录（默认 {DROP_DIR}）",
    )
    ap.add_argument(
        "--open",
        action="store_true",
        help="完成后用资源管理器打开输出目录",
    )
    args = ap.parse_args()
    drop: Path = args.drop

    print("=== 音师音色打包（USB / SoundWalkerIAP，非 J-Link）===")
    print(f"投放目录: {drop}")

    bin1_p = drop / "1.bin"
    bin2_p = drop / "2.bin"
    bin3_p = drop / "3.bin"
    missing = [p.name for p in (bin1_p, bin2_p, bin3_p) if not p.is_file()]
    if missing:
        raise SystemExit(
            f"投放目录缺少: {', '.join(missing)}\n"
            f"请把新的 1.bin / 2.bin / 3.bin 放到:\n  {drop}"
        )

    ziliao = find_ziliao()
    bin1 = bin1_p.read_bytes()
    bin2 = bin2_p.read_bytes()
    bin3 = bin3_p.read_bytes()
    logo = (ziliao / "logo.bin").read_bytes()
    charg = (ziliao / "Charg.bin").read_bytes()
    print(f"1.bin {len(bin1)} bytes | 2.bin {len(bin2)} bytes | 3.bin {len(bin3)} bytes")

    tone = pack_tone(bin1, bin2, bin3, logo, charg)
    ui = find_ui_baseline()
    all_res = tone + (b"\xFF" * (UI0902_RES_BASE - len(tone))) + ui
    require_equals(all_res, OFF_BIN1, bin1, "ALL/1.bin")
    require_equals(all_res, OFF_BIN2, bin2, "ALL/2.bin")
    require_equals(all_res, OFF_BIN3, bin3, "ALL/3.bin")
    dab_ok(all_res, OFF_BIN1, 31, "ALL/1.bin")
    dab_ok(all_res, OFF_BIN2, 1, "ALL/2.bin")
    dab_ok(all_res, OFF_BIN3, 3, "ALL/3.bin")

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    day = datetime.now().strftime("%Y%m%d")
    pkg = OUT_DIR / f"音师音色包_{day}"
    pkg.mkdir(parents=True, exist_ok=True)
    all_name = f"extflash_ALL_artist_{stamp}.res"
    all_path = pkg / all_name
    all_path.write_bytes(all_res)
    (OUT_DIR / "extflash_tone_artist.res").write_bytes(tone)
    (OUT_DIR / "extflash_ALL_artist.res").write_bytes(all_res)

    iap_exe = IAP_DIR / "SoundWalkerIAP.exe"
    iap_doc = IAP_DIR / "升级步骤.docx"
    if not iap_exe.is_file():
        raise SystemExit(f"找不到 {iap_exe}")
    shutil.copy2(iap_exe, pkg / "SoundWalkerIAP.exe")
    if iap_doc.is_file():
        shutil.copy2(iap_doc, pkg / "升级步骤.docx")
    write_flash_guide(pkg / "请这样刷.txt", all_name, stamp)

    print(f"OK 音色区 {len(tone)} bytes，ALL.res {len(all_res)} bytes")
    print(f"OK 输出目录: {pkg}")
    print(f"    - {all_name}")
    print("    - SoundWalkerIAP.exe")
    print("    - 请这样刷.txt")
    print("请用 USB + SoundWalkerIAP 刷 ALL.res，不要用 J-Link 单独烧 bin。")

    if args.open:
        try:
            import os

            os.startfile(str(pkg))  # type: ignore[attr-defined]
        except Exception as exc:
            print(f"(无法自动打开目录: {exc})")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except SystemExit as e:
        code = e.code
        if isinstance(code, str):
            print(f"\n失败: {code}")
            sys.exit(1)
        if code not in (0, None):
            sys.exit(int(code) if isinstance(code, int) else 1)
        sys.exit(0)
    except Exception as e:
        print(f"\n失败: {e}")
        sys.exit(1)