#!/usr/bin/env python3

import json
import tempfile
from pathlib import Path

import validate_partition_closure


def main() -> int:
    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        psp_json = root / "psp.json"
        table = [
            ("boot", "2560K"), ("env", "512K"), ("yoc", "1024K"),
            ("kernel.A", "3584K"), ("rootfs.A", "3072K"), ("firmware.A", "19456K"),
            ("kernel.B", "3584K"), ("rootfs.B", "3072K"), ("firmware.B", "19456K"),
            ("data", "14336K"), ("configs", "-"),
        ]
        psp_json.write_text(json.dumps({"partition": {"flashImgSize": "128M", "partitionTable": [{"name": name, "size": size} for name, size in table]}}), encoding="utf-8")
        _flash, rows = validate_partition_closure.partition_rows(psp_json)
        default = root / "Default"
        assignments = "\n".join(f"{row['boot_name']}_PART_OFFSET={row['offset_bytes']:#x}\n{row['boot_name']}_PART_SIZE={row['size_bytes']:#x}" for row in rows)
        mtdparts = ",".join(("-" if row["name"] == "configs" else hex(row["size_bytes"])) + f"({row['boot_name']})" for row in rows)
        default.write_text(assignments + f"\nMTD=mtdparts=spi0.0:{mtdparts}\n", encoding="utf-8")
        release = root / "release.py"
        release.write_text('run("./mkAiData.py 104 none")\n', encoding="utf-8")
        result = validate_partition_closure.validate(psp_json, default, release)
        assert result["passed"] is True, result

        default.write_text(default.read_text(encoding="utf-8").replace("DATA_PART_SIZE=0xe00000", "DATA_PART_SIZE=0x100000"), encoding="utf-8")
        stale = validate_partition_closure.validate(psp_json, default, release)
        assert stale["passed"] is False
        assert any("data_part_size" in error for error in stale["errors"])

        default.write_text(assignments + f"\nMTD=mtdparts=spi0.0:{mtdparts}\n", encoding="utf-8")
        release.write_text('run("./mkAiData.py 280 none")\n', encoding="utf-8")
        too_large = validate_partition_closure.validate(psp_json, default, release)
        assert too_large["passed"] is False
        assert any("超过 DATA" in error for error in too_large["errors"])

    print("partition closure tests passed")
    return 0


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