#!/usr/bin/env python3
"""Compare a fresh sampling rerun with matching formal-grid cells."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--formal-json", type=Path, required=True)
    parser.add_argument("--rerun-json", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(16 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def canonical_hash(value: Any) -> str:
    return hashlib.sha256(
        json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
        ).encode()
    ).hexdigest()


def run_index(
    payload: dict[str, Any],
) -> dict[tuple[str, int, str], tuple[dict[str, Any], dict[str, Any]]]:
    return {
        (
            source["id"],
            run["base_seed"],
            output["condition"],
        ): (run, output)
        for source in payload["sources"]
        for run in source["runs"]
        for output in run["outputs"]
    }


def main() -> None:
    args = parse_args()
    for path in (args.formal_json, args.rerun_json):
        if not path.is_file():
            raise FileNotFoundError(path)
    formal = json.loads(
        args.formal_json.read_text(encoding="utf-8")
    )
    rerun = json.loads(
        args.rerun_json.read_text(encoding="utf-8")
    )
    if formal["protocol_id"] != rerun["protocol_id"]:
        raise RuntimeError("protocol IDs differ")
    if formal["model"]["revision"] != rerun["model"]["revision"]:
        raise RuntimeError("model revisions differ")
    if (
        formal["generation_contract"]
        != rerun["generation_contract"]
    ):
        raise RuntimeError("generation contracts differ")
    formal_rows = run_index(formal)
    rerun_rows = run_index(rerun)
    missing = sorted(set(rerun_rows) - set(formal_rows))
    if missing:
        raise RuntimeError(
            f"{len(missing)} rerun cells are absent from formal grid"
        )

    rows = []
    for key in sorted(rerun_rows):
        formal_run, formal_output = formal_rows[key]
        rerun_run, rerun_output = rerun_rows[key]
        checks = {
            "run_seed_exact": (
                formal_run["run_seed"] == rerun_run["run_seed"]
            ),
            "prompt_hash_exact": (
                formal_output["prompt_token_ids_sha256"]
                == rerun_output["prompt_token_ids_sha256"]
            ),
            "generated_token_ids_exact": (
                formal_output["generated_token_ids"]
                == rerun_output["generated_token_ids"]
            ),
            "decoded_text_exact": (
                formal_output["text"] == rerun_output["text"]
            ),
            "eos_state_exact": (
                formal_output["hit_eos"]
                == rerun_output["hit_eos"]
            ),
            "truncation_state_exact": (
                formal_output["stopped_at_max_new_tokens"]
                == rerun_output["stopped_at_max_new_tokens"]
            ),
            "cpu_rng_pre_state_exact": (
                formal_run["rng_state_before"]["cpu_sha256"]
                == rerun_run["rng_state_before"]["cpu_sha256"]
            ),
            "cuda_rng_pre_state_exact": (
                formal_run["rng_state_before"][
                    "cuda_combined_sha256"
                ]
                == rerun_run["rng_state_before"][
                    "cuda_combined_sha256"
                ]
            ),
        }
        rows.append(
            {
                "source_id": key[0],
                "base_seed": key[1],
                "condition": key[2],
                **checks,
                "all_preregistered_fields_exact": all(
                    checks.values()
                ),
            }
        )
    check_names = [
        key
        for key in rows[0]
        if key.endswith("_exact")
        and key != "all_preregistered_fields_exact"
    ]
    result = {
        "schema_version": 1,
        "protocol_id": formal["protocol_id"],
        "formal": {
            "path": str(args.formal_json),
            "sha256": sha256_file(args.formal_json),
            "content_hash": formal["content_hash"],
            "base_seeds": formal["seed_contract"][
                "executed_base_seeds"
            ],
        },
        "rerun": {
            "path": str(args.rerun_json),
            "sha256": sha256_file(args.rerun_json),
            "content_hash": rerun["content_hash"],
            "base_seeds": rerun["seed_contract"][
                "executed_base_seeds"
            ],
        },
        "rows": rows,
        "summary": {
            "cells": len(rows),
            "all_preregistered_fields_exact": sum(
                row["all_preregistered_fields_exact"]
                for row in rows
            ),
            "by_field": {
                name: sum(row[name] for row in rows)
                for name in check_names
            },
        },
        "claim_boundary": [
            "Only the rerun seed subset is independently reproduced.",
            "Exact replay is scoped to the pinned software and hardware contract.",
            "Reproduction does not imply trajectories are seed-invariant.",
        ],
    }
    result["content_hash"] = canonical_hash(rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    payload = args.output.read_bytes()
    print(
        json.dumps(
            {
                "output": str(args.output),
                "sha256": hashlib.sha256(payload).hexdigest(),
                "bytes": len(payload),
                "summary": result["summary"],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()