from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any, Dict, List


def canonical_json(obj: Any) -> str:
    return json.dumps(obj, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def sha256_json(obj: Any) -> str:
    return sha256_text(canonical_json(obj))


def verify(path: Path) -> Dict[str, Any]:
    bundle = json.loads(path.read_text(encoding="utf-8"))
    payload = bundle.get("request_payload")
    expected_payload_hash = bundle.get("request_payload_sha256")
    actual_payload_hash = sha256_json(payload)

    trial_results: List[Dict[str, Any]] = []
    for trial in bundle.get("trials", []):
        if not trial.get("ok"):
            trial_results.append({"trial_index": trial.get("trial_index"), "ok": False, "reason": "trial_failed"})
            continue
        expected = trial.get("content_sha256")
        actual = sha256_text(trial.get("content", ""))
        trial_results.append(
            {
                "trial_index": trial.get("trial_index"),
                "ok": expected == actual,
                "expected_content_sha256": expected,
                "actual_content_sha256": actual,
            }
        )

    successful_hashes = [t.get("content_sha256") for t in bundle.get("trials", []) if t.get("ok")]
    fingerprints = [str(t.get("system_fingerprint")) for t in bundle.get("trials", []) if t.get("ok")]

    return {
        "path": str(path),
        "payload_hash_ok": expected_payload_hash == actual_payload_hash,
        "expected_payload_sha256": expected_payload_hash,
        "actual_payload_sha256": actual_payload_hash,
        "trial_content_hashes_ok": all(t["ok"] for t in trial_results),
        "all_successful_outputs_identical": len(set(successful_hashes)) <= 1 if successful_hashes else False,
        "all_successful_fingerprints_identical": len(set(fingerprints)) <= 1 if fingerprints else False,
        "unique_output_hashes": sorted(set(successful_hashes)),
        "unique_system_fingerprints": sorted(set(fingerprints)),
        "trials": trial_results,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify hashes inside LLM reproducibility run bundles.")
    parser.add_argument("bundle", nargs="+", type=Path)
    args = parser.parse_args()

    for path in args.bundle:
        result = verify(path)
        print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()