#!/usr/bin/env python3
"""Validate PSP partition geometry, Boot/Default mirrors, and DATA UBI capacity."""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import tempfile
from pathlib import Path


SIZE_RE = re.compile(r"^(\d+)([KMG])$", re.IGNORECASE)
MULTIPLIERS = {"K": 1024, "M": 1024**2, "G": 1024**3}
ALIASES = {
    "boot": "fip", "env": "ENV", "yoc": "2nd", "kernel.A": "BOOT",
    "rootfs.A": "ROOTFS", "firmware.A": "APP", "kernel.B": "BOOT2",
    "rootfs.B": "ROOTFS2", "firmware.B": "APP2", "data": "DATA", "configs": "CONFIGS",
}


def size_bytes(value: object) -> int | None:
    if value == "-":
        return None
    if not isinstance(value, str) or (match := SIZE_RE.fullmatch(value.strip())) is None:
        raise ValueError(f"非法分区大小: {value!r}")
    return int(match.group(1)) * MULTIPLIERS[match.group(2).upper()]


def parse_number(value: str) -> int:
    value = value.strip()
    match = SIZE_RE.fullmatch(value)
    if match:
        return int(match.group(1)) * MULTIPLIERS[match.group(2).upper()]
    return int(value, 0)


def partition_rows(psp_json: Path) -> tuple[int, list[dict[str, object]]]:
    payload = json.loads(psp_json.read_text(encoding="utf-8"))
    partition = payload.get("partition")
    if not isinstance(partition, dict):
        raise ValueError("psp.json.partition 必须是对象")
    flash = size_bytes(partition.get("flashImgSize"))
    if flash is None:
        raise ValueError("flashImgSize 不能是 -")
    table = partition.get("partitionTable")
    if not isinstance(table, list) or not table:
        raise ValueError("partitionTable 必须是非空数组")
    offset = 0
    rows: list[dict[str, object]] = []
    remainder_seen = False
    for index, item in enumerate(table):
        if not isinstance(item, dict) or not isinstance(item.get("name"), str):
            raise ValueError(f"partitionTable[{index}] 非法")
        raw_size = item.get("size")
        value = size_bytes(raw_size)
        if value is None:
            if remainder_seen or index != len(table) - 1:
                raise ValueError("size=- 只能出现一次且必须是最后一个分区")
            remainder_seen = True
            value = flash - offset
        if value < 0 or offset + value > flash:
            raise ValueError(f"分区 {item['name']} 超出 flashImgSize")
        rows.append({"name": item["name"], "boot_name": ALIASES.get(item["name"], item["name"]), "offset_bytes": offset, "size_bytes": value})
        offset += value
    if offset != flash:
        raise ValueError(f"分区总量 {offset} 未闭合到 flashImgSize {flash}")
    return flash, rows


def assignments(text: str) -> dict[str, str]:
    result: dict[str, str] = {}
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        match = re.match(r"(?:export\s+)?([A-Za-z0-9_]+)\s*=\s*['\"]?([^'\"\s;]+)", stripped)
        if match:
            result[match.group(1).casefold()] = match.group(2)
    return result


def check_mtdparts(
    text: str, rows: list[dict[str, object]], errors: list[str]
) -> dict[str, object]:
    candidates = sorted(
        set(re.findall(r"mtdparts=[^:\s\"']+:([^\"'\s;]+)", text))
    )
    if not candidates:
        return {"present": False, "checks": []}
    if len(candidates) != 1:
        errors.append("Boot/Default 中存在多个不同 mtdparts，无法唯一核对")
        return {"present": True, "values": candidates, "checks": []}
    tokens = candidates[0].split(",")
    checks: list[dict[str, object]] = []
    if len(tokens) != len(rows):
        errors.append(f"mtdparts 分区数量不一致：期望 {len(rows)}，实际 {len(tokens)}")
    for index, row in enumerate(rows):
        token = tokens[index] if index < len(tokens) else ""
        match = re.fullmatch(r"([^()]+)\(([^)]+)\)", token)
        raw_size, actual_name = (match.group(1), match.group(2)) if match else (None, None)
        try:
            actual_size = (
                int(row["size_bytes"])
                if raw_size == "-"
                else parse_number(raw_size) if raw_size is not None else None
            )
        except ValueError:
            actual_size = None
        expected_name = str(row["boot_name"])
        matched = (
            actual_name is not None
            and actual_name.casefold() == expected_name.casefold()
            and actual_size == row["size_bytes"]
        )
        checks.append({"index": index, "token": token, "expected_name": expected_name, "actual_name": actual_name, "expected_size": row["size_bytes"], "actual_size": actual_size, "matched": matched})
        if not matched:
            errors.append(f"mtdparts[{index}] 未同步：期望 {expected_name}/{row['size_bytes']}，实际 {token!r}")
    return {"present": True, "value": candidates[0], "checks": checks}


def check_embedded_partition_json(
    text: str, rows: list[dict[str, object]], errors: list[str]
) -> dict[str, object]:
    payloads: list[dict[str, object]] = []
    for line in text.splitlines():
        if "partitionTable" not in line or "{" not in line:
            continue
        normalized_line = line.replace('\\"', '"')
        start, end = normalized_line.find("{"), normalized_line.rfind("}")
        if start < 0 or end <= start:
            continue
        try:
            payload = json.loads(normalized_line[start : end + 1])
        except json.JSONDecodeError:
            errors.append("Boot/Default 内嵌 partition JSON 无法解析")
            continue
        if isinstance(payload, dict):
            payloads.append(payload)
    if not payloads:
        return {"present": False, "checks": []}
    if len(payloads) != 1:
        errors.append("Boot/Default 中存在多个内嵌 partition JSON，无法唯一核对")
        return {"present": True, "count": len(payloads), "checks": []}
    table = payloads[0].get("partitionTable")
    checks: list[dict[str, object]] = []
    if not isinstance(table, list) or len(table) != len(rows):
        errors.append("Boot/Default 内嵌 partition JSON 分区数量不一致")
        return {"present": True, "checks": checks}
    for index, row in enumerate(rows):
        item = table[index]
        actual_name = item.get("name") if isinstance(item, dict) else None
        try:
            actual_size = size_bytes(item.get("size")) if isinstance(item, dict) else None
            if actual_size is None:
                actual_size = row["size_bytes"]
        except ValueError:
            actual_size = None
        matched = actual_name == row["name"] and actual_size == row["size_bytes"]
        checks.append({"index": index, "expected_name": row["name"], "actual_name": actual_name, "expected_size": row["size_bytes"], "actual_size": actual_size, "matched": matched})
        if not matched:
            errors.append(f"内嵌 partition JSON[{index}] 未同步")
    return {"present": True, "checks": checks}


def validate(
    psp_json: Path,
    boot_default: Path,
    release_py: Path | None,
    leb_size: int = 126976,
    peb_size: int = 131072,
    reserve_pebs: int = 8,
) -> dict[str, object]:
    errors: list[str] = []
    psp_json = psp_json.expanduser().resolve(strict=True)
    boot_default = boot_default.expanduser().resolve(strict=True)
    if psp_json.is_symlink() or boot_default.is_symlink():
        raise ValueError("分区输入必须是普通文件，不能是软链接")
    flash, rows = partition_rows(psp_json)
    boot_text = boot_default.read_text(encoding="utf-8", errors="replace")
    boot_values = assignments(boot_text)
    mirror_checks: list[dict[str, object]] = []
    for row in rows:
        alias = str(row["boot_name"])
        keys = {
            "offset": f"{alias}_PART_OFFSET".casefold(),
            "size": f"{alias}_PART_SIZE".casefold(),
        }
        for kind, key in keys.items():
            expected = int(row[f"{kind}_bytes"])
            raw = boot_values.get(key)
            try:
                actual = parse_number(raw) if raw is not None else None
            except ValueError:
                actual = None
            matched = actual == expected
            mirror_checks.append({"partition": row["name"], "field": kind, "key": key, "expected": expected, "actual": actual, "matched": matched})
            if not matched:
                errors.append(f"Boot/Default {key} 未同步：期望 {expected:#x}，实际 {raw!r}")

    data = next((row for row in rows if row["name"] == "data"), None)
    mtdparts = check_mtdparts(boot_text, rows, errors)
    embedded_partition = check_embedded_partition_json(boot_text, rows, errors)
    data_check = None
    if release_py is not None:
        release_py = release_py.expanduser().resolve(strict=True)
        calls: set[int] = set()
        for line in release_py.read_text(encoding="utf-8", errors="replace").splitlines():
            if line.lstrip().startswith("#"):
                continue
            calls.update(int(value) for value in re.findall(r"mkAiData\.py\s+(\d+)\b", line))
        if len(calls) != 1:
            errors.append("release.py 必须唯一确定一个有效 mkAiData.py 容量参数")
        elif data is None:
            errors.append("PSP 分区表缺少 data")
        else:
            count = next(iter(calls))
            requested = count * leb_size
            reserved = reserve_pebs * peb_size
            capacity = int(data["size_bytes"])
            passed = requested + reserved <= capacity
            data_check = {"parameter": count, "leb_size": leb_size, "requested_bytes": requested, "reserve_pebs": reserve_pebs, "reserved_bytes": reserved, "partition_bytes": capacity, "passed": passed}
            if not passed:
                errors.append(f"mkAiData.py {count} 需要 {requested} 字节并保留 {reserved} 字节，超过 DATA {capacity} 字节")

    return {
        "schema_version": 1,
        "passed": not errors,
        "psp_json": str(psp_json),
        "boot_default": str(boot_default),
        "release_py": str(release_py) if release_py else None,
        "flash_size_bytes": flash,
        "partitions": rows,
        "boot_mirror_checks": mirror_checks,
        "mtdparts": mtdparts,
        "embedded_partition_json": embedded_partition,
        "data_capacity": data_check,
        "errors": errors,
    }


def write_json(path: Path, payload: dict[str, object]) -> None:
    path = path.expanduser().resolve(strict=False)
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--psp-json", type=Path, required=True)
    parser.add_argument("--boot-default", type=Path, required=True)
    parser.add_argument("--release-py", type=Path)
    parser.add_argument("--leb-size", type=int, default=126976)
    parser.add_argument("--peb-size", type=int, default=131072)
    parser.add_argument("--reserve-pebs", type=int, default=8)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    try:
        result = validate(args.psp_json, args.boot_default, args.release_py, args.leb_size, args.peb_size, args.reserve_pebs)
        if args.output:
            write_json(args.output, result)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0 if result["passed"] else 2
    except (OSError, UnicodeError, ValueError, json.JSONDecodeError) as exc:
        print(json.dumps({"passed": False, "error": str(exc)}, ensure_ascii=False, indent=2))
        return 2


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