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

流程（在 cargo build 之后）：

    package_firmware.py
        ├─ 从 Cargo 产物导出 App，同步 Bootloader / 分区表
        ├─ 刷新 BUILD-INFO.json，并调用 make_bundle.py 生成 bundle.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, timezone
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 resolve_elf(repo: Path, explicit: str) -> Path:
    if explicit:
        return Path(explicit).resolve()
    # Cargo metadata also respects build.target-dir in Cargo configuration.
    result = subprocess.run(
        ["cargo", "metadata", "--no-deps", "--format-version", "1"],
        cwd=repo, capture_output=True, text=True, check=True,
    )
    target = Path(json.loads(result.stdout)["target_directory"])
    return target / "xtensa-esp32s3-espidf" / "release" / "esp32"


def export_app(repo: Path, firmware: Path, elf: Path, allow_stale: bool) -> None:
    """从选中的 ELF 导出 App，并同步同次构建的配套镜像。"""
    app = firmware / "esp32s3-app.bin"
    if not elf.is_file():
        if allow_stale and app.is_file():
            print(f"警告：缺少 {elf}；继续打包现有镜像", file=sys.stderr)
            return
        raise SystemExit(f"固件 ELF 不存在：{elf}。请先构建，或用 --elf 指定本次构建的 ELF。")

    companions = [elf.parent / name for name in ("bootloader.bin", "partition-table.bin")]
    for source in companions:
        if not source.is_file():
            raise SystemExit(f"缺少同次构建的配套镜像：{source}")

    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}。请安装 espflash，或加 --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",
            "--flash-size",
            "16mb",
            "--flash-mode",
            "dio",
            "--flash-freq",
            "40mhz",
            "--skip-update-check",
            str(elf),
            str(app),
        ]
    )
    if result.returncode != 0:
        raise SystemExit(f"espflash save-image 失败（exit {result.returncode}）")

    for source in companions:
        destination = firmware / source.name
        if source.resolve() != destination.resolve():
            shutil.copy2(source, destination)

    # Do not attach a previous build's hardware/performance claims to new images.
    revision = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo, text=True).strip()
    dirty = subprocess.check_output(
        ["git", "status", "--porcelain", "--untracked-files=normal"], cwd=repo, text=True,
    ).strip()
    build = {
        "packaged_at": datetime.fromtimestamp(elf.stat().st_mtime, timezone.utc).isoformat(),
        "source_revision": revision,
        "includes_local_changes": bool(dirty),
        "target": "xtensa-esp32s3-espidf",
        "app_sha256": hashlib.sha256(app.read_bytes()).hexdigest(),
        "verification": {"hardware_flashed_and_tested": False},
    }
    (firmware / "BUILD-INFO.json").write_text(
        json.dumps(build, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="生成官方 .zip 固件包")
    parser.add_argument("--firmware", default="target/firmware", help="固件产物目录")
    parser.add_argument("--elf", default="", help="本次构建的 ELF；默认通过 cargo metadata 查找 release 产物")
    parser.add_argument("--name", default="", help="传递给 make_bundle.py 的名称")
    parser.add_argument("--version", default="", help="传递给 make_bundle.py 的版本；默认读取 version.json")
    parser.add_argument("--skip-manifest", action="store_true", help="不重新生成 bundle.json")
    parser.add_argument("--allow-stale", action="store_true", help="缺少 ELF 或 espflash 时，允许打包目录中的现有镜像")
    parser.add_argument("--output", default="", help="输出 zip 路径；默认 target/firmware/camfw-*.zip")
    args = parser.parse_args()

    repo = Path(__file__).resolve().parent.parent
    if not args.version and (repo / "version.json").is_file():
        version = json.loads((repo / "version.json").read_text(encoding="utf-8"))
        args.version = f"v{version['version']}+{version['build']}"
    firmware = Path(args.firmware)
    if not firmware.is_absolute():
        firmware = (Path.cwd() / firmware).resolve()
    try:
        export_app(repo, firmware, resolve_elf(repo, args.elf), args.allow_stale)
    except (SystemExit, OSError, subprocess.CalledProcessError, ValueError) 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())