#!/usr/bin/env python3
"""每次构建后报告分区和镜像空间占用。"""

from __future__ import annotations

import argparse
import struct
from pathlib import Path


ENTRY = struct.Struct("<2sBBII16sI")


def partitions(path: Path) -> list[tuple[str, int, int]]:
    # 读取本次构建生成的分区表二进制条目；报告不代表设备当前烧录状态。
    data = path.read_bytes()
    result = []
    for offset in range(0, len(data) - ENTRY.size + 1, ENTRY.size):
        magic, _type, _subtype, start, size, raw_name, _flags = ENTRY.unpack_from(data, offset)
        if magic != b"\xaa\x50":
            break
        result.append((raw_name.split(b"\0", 1)[0].decode("ascii"), start, size))
    if not result:
        raise ValueError(f"no partitions found in {path}")
    return result


def mib(size: int) -> str:
    return f"{size / 1048576:.2f} MiB"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--partition-table", required=True, type=Path)
    parser.add_argument("--app", required=True, type=Path)
    parser.add_argument("--asset-pack", required=True, type=Path)
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--flash-size", required=True)
    args = parser.parse_args()

    flash_size = int(args.flash_size.upper().removesuffix("MB")) * 1048576
    entries = partitions(args.partition_table)
    by_name = {name: (start, size) for name, start, size in entries}
    # 素材包已计入应用镜像，单独显示仅用于分析，不能再加到应用占用中。
    app_size = args.app.stat().st_size
    asset_size = args.asset_pack.stat().st_size
    model_size = args.model.stat().st_size
    for name in ("ota_0", "ota_1", "model"):
        if name not in by_name:
            raise ValueError(f"missing required partition: {name}")
    if app_size > min(by_name[name][1] for name in ("ota_0", "ota_1")):
        raise ValueError("application image exceeds an OTA slot")
    if model_size > by_name["model"][1]:
        raise ValueError("model image exceeds model partition")
    last_end = max(start + size for _, start, size in entries)
    if last_end > flash_size:
        raise ValueError("partition table exceeds configured Flash size")

    print("\nFlash usage (local build images; device contents may differ):")
    print(f"  Flash capacity: {mib(flash_size)}")
    for name, start, size in entries:
        print(f"  {name:9s} 0x{start:06x}-0x{start + size:06x} {mib(size)}")
    slot_size = min(by_name[name][1] for name in ("ota_0", "ota_1"))
    print(f"  App image:     {app_size:,} B / {slot_size:,} B ({app_size / slot_size:.1%}); free {mib(slot_size - app_size)} per OTA slot")
    print(f"  Embedded assets: {asset_size:,} B (included in app image)")
    print(f"  Model image:   {model_size:,} B / {by_name['model'][1]:,} B ({model_size / by_name['model'][1]:.1%})")
    print(f"  Unpartitioned tail: {mib(flash_size - last_end)}\n")


if __name__ == "__main__":
    main()