#!/usr/bin/env python3
"""Verify that frozen pair/action manifests resolve against reconstructed data."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path


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


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--processed-root", type=Path, required=True)
    parser.add_argument("--checksums", type=Path)
    parser.add_argument("--action-manifest-dir", type=Path)
    parser.add_argument("--derived-root", type=Path)
    args = parser.parse_args()

    rows = [json.loads(line) for line in args.manifest.read_text().splitlines() if line]
    references: set[Path] = set()
    for row in rows:
        for role in ("chosen", "rejected"):
            relative = Path(str(row[role]["source_frame_file"]))
            parts = relative.parts[1:] if relative.parts[:1] in ((".",), ("processed_datasets",)) else relative.parts
            if relative.parts[:1] == (".",) and parts[:1] == ("processed_datasets",):
                parts = parts[1:]
            references.add(Path(*parts))
    missing = [str(path) for path in sorted(references) if not (args.processed_root / path).is_file()]
    if missing:
        raise FileNotFoundError(f"missing {len(missing)} manifest frame files; examples={missing[:10]}")
    if len(rows) != 2689 or len(references) != 4323:
        raise RuntimeError(f"unexpected frozen counts: rows={len(rows)} files={len(references)}")

    checksum_count = 0
    if args.checksums:
        for line in args.checksums.read_text().splitlines():
            expected, relative = line.split(maxsplit=1)
            relative_path = Path(relative.strip())
            parts = relative_path.parts[1:] if relative_path.parts[:1] == ("processed_datasets",) else relative_path.parts
            target = args.processed_root / Path(*parts)
            if digest(target) != expected:
                raise RuntimeError(f"SHA256 mismatch: {target}")
            checksum_count += 1

    action_rows = action_files = 0
    if args.action_manifest_dir or args.derived_root:
        if not (args.action_manifest_dir and args.derived_root):
            raise ValueError("action-manifest-dir and derived-root must be supplied together")
        for manifest in sorted(args.action_manifest_dir.glob("*.jsonl")):
            for line in manifest.read_text().splitlines():
                row = json.loads(line)
                action_rows += 1
                if "sidecar_path" not in row:
                    continue
                target = args.derived_root / row["sidecar_path"]
                if not target.is_file():
                    raise FileNotFoundError(target)
                action_files += 1

    print(json.dumps({
        "status": "pass",
        "pair_rows": len(rows),
        "unique_stage2_frame_files": len(references),
        "sha256_files_verified": checksum_count,
        "action_manifest_rows": action_rows,
        "action_sidecar_files_resolved": action_files,
    }, indent=2))


if __name__ == "__main__":
    main()