#!/usr/bin/env python3
"""Validate external model or third-party asset manifests without loading them."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path
from typing import Any, Iterable


def _digest(path: Path) -> str:
    hasher = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            hasher.update(chunk)
    return hasher.hexdigest()


def _model_entries(manifest: dict[str, Any], backend: str | None) -> Iterable[dict[str, Any]]:
    yield from manifest.get("common", [])
    if backend:
        try:
            yield from manifest["backends"][backend]
        except KeyError as error:
            available = ", ".join(sorted(manifest.get("backends", {})))
            raise ValueError(f"Unknown backend {backend!r}; expected one of: {available}") from error
    else:
        for entries in manifest.get("backends", {}).values():
            yield from entries


def _third_party_entries(manifest: dict[str, Any]) -> Iterable[dict[str, Any]]:
    for entry in manifest.get("assets", []):
        yield entry


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", required=True, type=Path)
    parser.add_argument("--root", required=True, type=Path)
    parser.add_argument("--backend", help="validate one model backend instead of every model asset")
    parser.add_argument(
        "--verify-sha256",
        action="store_true",
        help="also calculate and compare SHA-256 values (optional, potentially slow)",
    )
    parser.add_argument("--verify-size", action="store_true", help="also compare recorded byte sizes")
    args = parser.parse_args()

    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    if "backends" in manifest:
        entries = list(_model_entries(manifest, args.backend))
    else:
        if args.backend:
            parser.error("--backend is only valid for a model manifest")
        entries = list(_third_party_entries(manifest))

    errors: list[str] = []
    for entry in entries:
        path = args.root / entry["path" if "path" in entry else "target"]
        label = entry.get("role", entry.get("id", path.name))
        if not path.is_file():
            errors.append(f"missing {label}: {path}")
            continue
        expected_size = entry.get("size_bytes")
        if args.verify_size and expected_size is not None and path.stat().st_size != expected_size:
            errors.append(f"wrong size for {label}: {path} (expected {expected_size}, got {path.stat().st_size})")
            continue
        if args.verify_sha256 and entry.get("sha256") and _digest(path) != entry["sha256"]:
            errors.append(f"SHA-256 mismatch for {label}: {path}")
    if errors:
        print("Asset verification failed:", file=sys.stderr)
        print("\n".join(f"- {error}" for error in errors), file=sys.stderr)
        return 1
    print(f"Verified {len(entries)} asset(s) from {args.manifest} against {args.root}")
    return 0


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