#!/usr/bin/env python3
"""check_determinism.py — Layer 0: verify reproducibility chain.

Checks:
  1. Every artifact in extraction/ has a content_sha256 field that matches a
     re-computation of the canonical JSON (sorted keys, no whitespace).
  2. No Python shim calls unseeded RNG (np.random.default_rng() with no seed).
  3. All --seed parameters default to 0.

Exit codes:
  0 = All hash chains match (deterministic)
  1 = Content hash mismatch
  2 = Missing receipt or source file
  3 = Seed-lock violation detected
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import sys
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parent.parent


def compute_file_sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def compute_json_sha256(data: dict | list) -> str:
    """Compute SHA-256 of canonical JSON, excluding self-referential content_sha256 field."""
    if isinstance(data, dict):
        data = {k: v for k, v in data.items() if k not in ("content_sha256", "receipt_hash")}
    canonical = json.dumps(data, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode()).hexdigest()


def check_artifact_chain(receipt_dir: Path) -> dict[str, Any]:
    """Verify SHA-256 chain for all JSON artifacts in receipt_dir."""
    results = {"checked": 0, "passed": 0, "failed": 0, "missing": 0, "artifacts": []}

    for path in sorted(receipt_dir.glob("*.json")):
        results["checked"] += 1
        try:
            data = json.loads(path.read_text())
        except (json.JSONDecodeError, IOError):
            results["failed"] += 1
            results["artifacts"].append({"path": str(path), "status": "parse_error"})
            continue

        stored_hash = data.get("content_sha256") or data.get("receipt", {}).get("receipt_hash")
        if not stored_hash:
            results["missing"] += 1
            results["artifacts"].append({"path": str(path), "status": "no_hash"})
            continue

        computed = compute_json_sha256(data)
        if computed == stored_hash:
            results["passed"] += 1
            results["artifacts"].append({"path": str(path), "status": "match"})
        else:
            results["failed"] += 1
            results["artifacts"].append({
                "path": str(path), "status": "mismatch",
                "stored": stored_hash[:16],
                "computed": computed[:16],
            })

    return results


def scan_seed_violations(root_dirs: list[Path]) -> list[dict[str, Any]]:
    """Find unseeded RNG calls in Python shims."""
    violations: list[dict[str, Any]] = []
    pattern = re.compile(r"np\.random\.default_rng\(\s*\)")

    for root in root_dirs:
        for py_file in root.rglob("*.py"):
            if ".lake" in str(py_file) or "__pycache__" in str(py_file):
                continue
            text = py_file.read_text()
            # Check for unseeded numpy RNG
            for match in pattern.finditer(text):
                violations.append({
                    "file": str(py_file.relative_to(REPO_ROOT)),
                    "line": text[:match.start()].count("\n") + 1,
                    "code": match.group(),
                    "fix": "np.random.default_rng(seed)  # where seed comes from --seed arg",
                })
    return violations


def check_shim_seed_params(root_dirs: list[Path]) -> list[dict[str, Any]]:
    """Verify all CLI entry points accept --seed parameter."""
    issues: list[dict[str, Any]] = []
    seed_arg_pattern = re.compile(r'"--seed"')

    for root in root_dirs:
        for py_file in root.rglob("*.py"):
            if ".lake" in str(py_file) or "__pycache__" in str(py_file):
                continue
            text = py_file.read_text()
            # Only check files that have argparse
            if "argparse" not in text and "ArgumentParser" not in text:
                continue
            # Check if they have random/numpy RNG usage but no --seed
            has_rng = "np.random" in text or "random.Random" in text
            has_seed_arg = bool(seed_arg_pattern.search(text))
            if has_rng and not has_seed_arg:
                issues.append({
                    "file": str(py_file.relative_to(REPO_ROOT)),
                    "has_rng": True,
                    "has_seed_arg": False,
                    "fix": "Add parser.add_argument('--seed', type=int, default=0)",
                })
    return issues


def main():
    parser = argparse.ArgumentParser(description="Layer 0: Deterministic Reproducibility Chain")
    parser.add_argument("--seed", type=int, default=0, help="Seed for verification (default: 0)")
    parser.add_argument("--receipt-dir", type=Path, 
                        default=REPO_ROOT / "extraction",
                        help="Directory containing JSON artifacts")
    parser.add_argument("--check-all", action="store_true", help="Run all checks")
    parser.add_argument("--output", type=Path, default=None, help="Output receipt path")
    args = parser.parse_args()

    exit_code = 0

    # Check 1: Hash chain integrity
    print(f"[check] Verifying hash chain in {args.receipt_dir}")
    chain = check_artifact_chain(args.receipt_dir)
    if chain["failed"] > 0:
        print(f"  FAILED: {chain['failed']}/{chain['checked']} artifacts have hash mismatches")
        exit_code = 1
    else:
        print(f"  PASSED: {chain['passed']}/{chain['checked']} artifacts verified")
    if chain["missing"] > 0:
        print(f"  WARNING: {chain['missing']} artifacts have no content_sha256 field")

    # Check 2: Seed violations
    print(f"[check] Scanning for unseeded RNG calls")
    scan_dirs = [REPO_ROOT / "python", REPO_ROOT / "qubo"]
    violations = scan_seed_violations([d for d in scan_dirs if d.exists()])
    if violations:
        print(f"  FAILED: {len(violations)} unseeded RNG calls found")
        for v in violations[:5]:
            print(f"    {v['file']}:{v['line']}  {v['code']}")
        exit_code = 3
    else:
        print(f"  PASSED: No unseeded RNG calls")

    # Check 3: Shims missing --seed
    print(f"[check] Scanning for shims missing --seed parameter")
    issues = check_shim_seed_params([d for d in scan_dirs if d.exists()])
    if issues:
        print(f"  WARNING: {len(issues)} shims use RNG but lack --seed")
        for i in issues:
            print(f"    {i['file']}")
    else:
        print(f"  PASSED: All RNG-using shims have --seed")

    # Generate output receipt if requested
    if args.output:
        receipt = {
            "schema": "anti_smuggle_layer0_receipt_v1",
            "seed": args.seed,
            "hash_chain": chain,
            "seed_violations": len(violations),
            "missing_seed_params": len(issues),
            "exit_code": exit_code,
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(receipt, indent=2))
        print(f"[check] Receipt written to {args.output}")

    sys.exit(exit_code)


if __name__ == "__main__":
    main()
