#!/usr/bin/env python3
"""Check the outer distribution file set and hashes; never alter model files."""
import argparse
import hashlib
import json
from pathlib import Path

NAME = "DISTRIBUTION_MANIFEST.json"
MANIFEST_SHA256 = "eef92679b27ed4fac41873d7aafb4f67371de4f875b74c023db9fa50d5975b3a"


def digest(path):
    h = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def check(root, expected_outer_sha256):
    root = root.resolve(strict=True)
    outer = root / NAME
    if outer.is_symlink() or not outer.is_file():
        raise ValueError("missing or linked outer manifest")
    if digest(outer) != expected_outer_sha256:
        raise ValueError("outer manifest SHA-256 differs from the separately trusted value")
    spec = json.loads(outer.read_text(encoding="utf-8"))
    if spec.get("schema") != 1 or spec.get("original_checkpoint_manifest_sha256") != MANIFEST_SHA256:
        raise ValueError("wrong package schema or original checkpoint identity")
    files = spec.get("files")
    if not isinstance(files, dict) or not files or NAME in files:
        raise ValueError("invalid file map or self-entry")
    if not any(key.startswith("checkpoint/") for key in files):
        raise ValueError("checkpoint file set is absent")
    for name, item in files.items():
        p = Path(name)
        if (not isinstance(name, str) or p.is_absolute() or p.as_posix() != name
                or ".." in p.parts or "\\" in name or not isinstance(item, dict)
                or not isinstance(item.get("sha256"), str) or len(item["sha256"]) != 64
                or type(item.get("size")) is not int or item["size"] < 0):
            raise ValueError("unsafe or malformed file map entry")
    actual = set()
    for p in root.rglob("*"):
        rel = p.relative_to(root).as_posix()
        # Hugging Face local_dir metadata belongs to the OUTER root, never checkpoint/.
        if rel == ".cache" or rel.startswith(".cache/"):
            if rel == ".cache" and p.is_symlink():
                raise ValueError("linked outer cache root")
            if rel != ".cache" and not rel.startswith(".cache/huggingface/") and rel != ".cache/huggingface":
                raise ValueError("unexpected outer cache entry: " + rel)
            continue
        if p.is_symlink():
            raise ValueError("symlink outside allowed outer cache: " + rel)
        if p.is_file():
            actual.add(rel)
        elif not p.is_dir():
            raise ValueError("unsupported entry: " + rel)
    if actual != set(files) | {NAME}:
        raise ValueError("file-set mismatch: missing=%s extra=%s" % (sorted(set(files) - actual), sorted(actual - set(files) - {NAME})))
    checkpoint = root / "checkpoint"
    if not checkpoint.is_dir() or checkpoint.is_symlink():
        raise ValueError("missing or linked checkpoint directory")
    for name, item in files.items():
        p = root / name
        if p.stat().st_size != item["size"] or digest(p) != item["sha256"]:
            raise ValueError("size or SHA-256 mismatch: " + name)
    if files["checkpoint/axl-ablit-manifest.json"]["sha256"] != MANIFEST_SHA256:
        raise ValueError("original checkpoint manifest changed")
    print("PASS: outer file set and all inventory SHA-256 hashes (%d files); runtime contract not checked" % len(files))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("root", type=Path, help="outer download directory")
    parser.add_argument("--outer-sha256", required=True, help="trusted 64-character published digest of outer manifest")
    args = parser.parse_args()
    if len(args.outer_sha256) != 64 or any(c not in "0123456789abcdef" for c in args.outer_sha256):
        parser.error("--outer-sha256 must be a lowercase SHA-256")
    try:
        check(args.root, args.outer_sha256)
    except (OSError, ValueError, KeyError, json.JSONDecodeError) as e:
        parser.exit(1, "FAIL: %s\n" % e)