#!/usr/bin/env python3
"""把编译后的固件目录打成客户端可直接读取的 .zip 官方包。

流程（在 cargo build 之后）：

    make_bundle.py  →  bundle.json（沿用现有清单）
    package_firmware.py
        ├─ 读 bundle.json / BUILD-INFO.json
        ├─ 固定条目顺序与 mtime，生成可复现的 zip
        └─ 写 <zip>.sha256

zip 根目录结构：

    bundle.json
    esp32s3-app.bin
    partition-table.bin   (0xC00 字节)
    bootloader.bin         (可选)
    BUILD-INFO.json        (可选)

客户端读取 zip 时不要求先解压；目录包仍然兼容。
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
import zipfile
from datetime import datetime
from pathlib import Path


def sanitize(value: str) -> str:
    value = re.sub(r"[^A-Za-z0-9._-]+", "-", value.strip())
    return value.strip("-") or "firmware"


def parse_epoch(build: dict) -> int:
    raw = os.environ.get("SOURCE_DATE_EPOCH")
    if raw:
        try:
            return int(raw)
        except ValueError:
            pass

    packaged = build.get("packaged_at")
    if packaged:
        try:
            return int(datetime.fromisoformat(packaged).timestamp())
        except (ValueError, TypeError):
            pass

    return int(time.time())


def export_app_if_stale(repo: Path, firmware: Path, allow_stale: bool) -> None:
    """编译后 App 镜像可能还没导出；优先用 espflash 从最新 ELF 重新导出。"""
    elf = repo / "target" / "xtensa-esp32s3-espidf" / "release" / "esp32"
    app = firmware / "esp32s3-app.bin"
    if not elf.is_file():
        return
    if app.is_file() and app.stat().st_mtime >= elf.stat().st_mtime:
        return

    if app.is_file():
        print(f"提示：{app} 早于 {elf}，将尝试从最新 ELF 重新导出", file=sys.stderr)
    else:
        print(f"提示：缺少 {app}，将尝试从最新 ELF 导出", file=sys.stderr)

    espflash = repo / "target" / "tools" / "espflash" / "espflash.exe"
    if not espflash.is_file():
        found = shutil.which("espflash")
        espflash = Path(found) if found else Path("")
    if not espflash.is_file():
        message = "找不到 espflash，无法从最新 ELF 导出 App 镜像"
        if allow_stale:
            print(f"警告：{message}；继续打包现有镜像", file=sys.stderr)
            return
        raise SystemExit(f"{message}。请先导出 target/firmware/esp32s3-app.bin，或加 --allow-stale。")

    firmware.mkdir(parents=True, exist_ok=True)
    print(f"导出 App：{elf} -> {app}")
    result = subprocess.run(
        [
            str(espflash),
            "save-image",
            "--chip",
            "esp32s3",
            "--format",
            "esp-idf",
            str(elf),
            str(app),
        ]
    )
    if result.returncode != 0:
        raise SystemExit(f"espflash save-image 失败（exit {result.returncode}）")


def main() -> int:
    parser = argparse.ArgumentParser(description="生成官方 .zip 固件包")
    parser.add_argument("--firmware", default="target/firmware", help="固件产物目录")
    parser.add_argument("--name", default="", help="传递给 make_bundle.py 的名称")
    parser.add_argument("--version", default="", help="传递给 make_bundle.py 的版本")
    parser.add_argument("--skip-manifest", action="store_true", help="不重新生成 bundle.json")
    parser.add_argument("--allow-stale", action="store_true", help="允许打包早于 ELF 的 App 镜像")
    parser.add_argument("--output", default="", help="输出 zip 路径；默认 target/firmware/camfw-*.zip")
    args = parser.parse_args()

    repo = Path(__file__).resolve().parent.parent
    firmware = Path(args.firmware)
    if not firmware.is_absolute():
        firmware = (Path.cwd() / firmware).resolve()
    if not firmware.is_dir():
        print(f"固件目录不存在：{firmware}", file=sys.stderr)
        return 1

    try:
        export_app_if_stale(repo, firmware, args.allow_stale)
    except SystemExit as error:
        print(error, file=sys.stderr)
        return 1

    if not args.skip_manifest:
        make_bundle = repo / "scripts" / "make_bundle.py"
        command = [sys.executable, str(make_bundle), "--firmware", str(firmware)]
        if args.name:
            command += ["--name", args.name]
        if args.version:
            command += ["--version", args.version]
        result = subprocess.run(command, text=True)
        if result.returncode != 0:
            print("make_bundle.py 失败，未生成 zip", file=sys.stderr)
            return result.returncode

    bundle_path = firmware / "bundle.json"
    if not bundle_path.is_file():
        print(f"缺少 {bundle_path}", file=sys.stderr)
        return 1

    bundle = json.loads(bundle_path.read_text(encoding="utf-8"))
    build = {}
    build_info_path = firmware / "BUILD-INFO.json"
    if build_info_path.is_file():
        try:
            build = json.loads(build_info_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            build = {}

    app = firmware / bundle["app"]
    partition_table = firmware / bundle["partitionTable"]
    if not app.is_file() or not partition_table.is_file():
        print("bundle.json 里的 App/分区表文件不存在", file=sys.stderr)
        return 1
    if partition_table.stat().st_size != 0xC00:
        print("partition-table.bin 必须是 0xC00 字节", file=sys.stderr)
        return 1

    bootloader_name = bundle.get("bootloader")
    bootloader = firmware / bootloader_name if bootloader_name else None
    if bootloader is not None and not bootloader.is_file():
        print(f"bundle.json 声明的 bootloader 不存在：{bootloader}", file=sys.stderr)
        return 1

    epoch = parse_epoch(build)
    stamp = time.gmtime(epoch)
    date_time = (
        stamp.tm_year,
        stamp.tm_mon,
        stamp.tm_mday,
        stamp.tm_hour,
        stamp.tm_min,
        stamp.tm_sec,
    )
    date_text = f"{stamp.tm_year:04d}{stamp.tm_mon:02d}{stamp.tm_mday:02d}"

    variant = sanitize(str(build.get("build_variant", "firmware")))
    revision = sanitize(str(build.get("source_revision", ""))[:7])
    chip = sanitize(str(bundle.get("chip", "esp32s3")))

    # 文件名优先用**传进来的版本**（`--version`，来自仓库根的 version.json）。
    #
    # 原来只用 BUILD-INFO.json 里的变体名 + 提交号 + 日期，那是另一套编号：固件包叫
    # `...-4794bed-20261018.zip`，而客户端里显示的是 `v0.2.0+34`，两个数对不上，
    # 现场拿到一个 zip 没法说清它是不是界面上那一版。给了版本就用版本，没给才退回旧格式。
    if args.version:
        filename = f"camfw-{chip}-{sanitize(args.version)}.zip"
    else:
        suffix = f"-{revision}" if revision else ""
        filename = f"camfw-{chip}-{variant}{suffix}-{date_text}.zip"

    output = Path(args.output) if args.output else firmware / filename
    if not output.is_absolute():
        output = (Path.cwd() / output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    entries: list[tuple[str, bytes]] = [
        ("bundle.json", bundle_path.read_bytes()),
        (app.name, app.read_bytes()),
        (partition_table.name, partition_table.read_bytes()),
    ]
    if bootloader is not None:
        entries.append((bootloader.name, bootloader.read_bytes()))
    if build_info_path.is_file():
        entries.append(("BUILD-INFO.json", build_info_path.read_bytes()))

    with zipfile.ZipFile(
        output,
        mode="w",
        compression=zipfile.ZIP_DEFLATED,
        compresslevel=9,
    ) as archive:
        for name, data in entries:
            info = zipfile.ZipInfo(filename=name, date_time=date_time)
            info.compress_type = zipfile.ZIP_DEFLATED
            info.create_system = 3
            info.external_attr = 0o100644 << 16
            archive.writestr(info, data)

    digest = hashlib.sha256(output.read_bytes()).hexdigest()
    sidecar = output.with_name(output.name + ".sha256")
    sidecar.write_text(f"{digest}  {output.name}\n", encoding="utf-8")

    print(f"已生成 {output}")
    print(f"  {output.stat().st_size} 字节 · {len(entries)} 项 · SHA256 {digest}")
    print(f"  {sidecar}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())