#!/usr/bin/env python3
"""将素材确定性地打包为只读资源，随应用镜像一起参与 OTA。"""

from __future__ import annotations

import argparse
import struct
from pathlib import Path


MAGIC = b"TQASSET1"
# 与 platform_asset_bundle.c 的解析格式一致：小端包头及条目头，后接名称和原始字节。
HEADER = struct.Struct("<8sI")
ENTRY = struct.Struct("<HI")
ALLOWED = {"lcd": ".png", "audio_prompts": ".opus"}


def pack(source: Path) -> bytes:
    # 排序固定条目顺序，保证相同输入生成相同的包；只接受当前运行时使用的两类素材。
    files = sorted(path for path in source.rglob("*") if path.is_file())
    if not files:
        raise ValueError("asset staging directory is empty")

    records: list[tuple[bytes, bytes]] = []
    names: set[str] = set()
    for path in files:
        relative = path.relative_to(source)
        if len(relative.parts) != 2 or relative.parts[0] not in ALLOWED:
            raise ValueError(f"unexpected asset path: {relative}")
        if path.suffix.lower() != ALLOWED[relative.parts[0]]:
            raise ValueError(f"unexpected asset type: {relative}")
        name = relative.as_posix()
        if name in names:
            raise ValueError(f"duplicate asset: {name}")
        names.add(name)
        encoded = name.encode("utf-8")
        payload = path.read_bytes()
        if not encoded or len(encoded) > 0xFFFF or not payload or len(payload) > 0xFFFFFFFF:
            raise ValueError(f"invalid asset size: {name}")
        records.append((encoded, payload))

    output = bytearray(HEADER.pack(MAGIC, len(records)))
    for name, payload in records:
        output.extend(ENTRY.pack(len(name), len(payload)))
        output.extend(name)
        output.extend(payload)
    print(f"Embedded assets: {len(records)} files, {sum(len(data) for _, data in records)} payload bytes, {len(output)} bundle bytes")
    return bytes(output)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    data = pack(args.source)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if not args.output.exists() or args.output.read_bytes() != data:
        args.output.write_bytes(data)


if __name__ == "__main__":
    main()