#!/usr/bin/env python3
"""cross_validate.py — Layer 1: Multi-model cross-validation.

Takes Lean proof files from two different models, extracts theorem
signatures, builds both independently, and checks equivalence.

CLI:
  python3 scripts/cross_validate.py \
    --model-a path/to/A.lean --model-b path/to/B.lean \
    --label-a gemma4-12b --label-b deepseek-v4-flash

Exit codes:
  0 = All common theorems equivalent
  1 = Non-equivalent theorems found
  2 = Build failure (one or both files don't compile)
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

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

# Regex: extract theorem name and signature from Lean source
THEOREM_RE = re.compile(
    r"(?:theorem|lemma)\s+(?P<name>\w+)\s*(?P<signature>.*?)(?::=)",
    re.DOTALL,
)


def extract_theorems(lean_text: str) -> dict[str, str]:
    """Extract {theorem_name: statement_type} from Lean source."""
    theorems: dict[str, str] = {}
    for match in THEOREM_RE.finditer(lean_text):
        name = match.group("name")
        sig = match.group("signature").strip()
        # Normalize whitespace
        sig = " ".join(sig.split())
        theorems[name] = sig
    return theorems


def normalize_statement(stmt: str) -> str:
    """Normalize binder names for comparison: binder_0, binder_1, ..."""
    # Replace variable-like names (single lowercase letters) with canonical forms
    import re
    var_pattern = re.compile(r'\b([a-z])\b')
    seen: dict[str, str] = {}
    counter = 0

    def _replace(m):
        nonlocal counter
        v = m.group(1)
        if v not in seen:
            seen[v] = f"x{counter}"
            counter += 1
        return seen[v]

    return var_pattern.sub(_replace, stmt)


def statements_equivalent(sig_a: str, sig_b: str) -> bool:
    """Check two theorem signatures for alpha-equivalence."""
    return normalize_statement(sig_a) == normalize_statement(sig_b)


def run_build(lean_file: Path, target: str = "SilverSightRRC",
              timeout_s: int = 120) -> tuple[int, str]:
    """Run lake build and return (exit_code, output)."""
    # Temporarily swap the file
    bak = lean_file.with_suffix(".lean.bak")
    if not bak.exists():
        lean_file.rename(bak)

    # Build
    try:
        result = subprocess.run(
            ["lake", "build", target],
            cwd=REPO_ROOT, capture_output=True, text=True, timeout=timeout_s,
        )
        return result.returncode, result.stdout + result.stderr
    except subprocess.TimeoutExpired:
        return -1, "TIMEOUT"
    finally:
        if bak.exists():
            bak.rename(lean_file)


def build_independent(lean_file: Path, target: str) -> tuple[int, str, str]:
    """Build the full target and record the Lean file's hash."""
    content = lean_file.read_bytes()
    sha = hashlib.sha256(content).hexdigest()
    try:
        result = subprocess.run(
            ["lake", "build", target],
            cwd=REPO_ROOT, capture_output=True, text=True, timeout=120,
        )
        return result.returncode, result.stdout + result.stderr, sha
    except subprocess.TimeoutExpired:
        return -1, "TIMEOUT", sha


def cross_validate(file_a: Path, file_b: Path, target: str,
                   label_a: str = "model_a", label_b: str = "model_b",
                   timeout_s: int = 120) -> dict[str, Any]:
    """Full cross-validation pipeline."""
    text_a = file_a.read_text()
    text_b = file_b.read_text()

    # Extract theorem signatures
    theorems_a = extract_theorems(text_a)
    theorems_b = extract_theorems(text_b)
    common = set(theorems_a.keys()) & set(theorems_b.keys())
    only_a = set(theorems_a.keys()) - set(theorems_b.keys())
    only_b = set(theorems_b.keys()) - set(theorems_a.keys())

    # Check equivalence for common theorems
    equiv_results = []
    all_equivalent = True
    for name in sorted(common):
        eq = statements_equivalent(theorems_a[name], theorems_b[name])
        if not eq:
            all_equivalent = False
        equiv_results.append({
            "name": name,
            "sig_a": theorems_a[name][:100],
            "sig_b": theorems_b[name][:100],
            "equivalent": eq,
        })

    # Build both independently
    print(f"[crossval] Building {label_a} ({file_a.name})...")
    exit_a, output_a, hash_a = build_independent(file_a, target)
    build_a_ok = exit_a == 0

    print(f"[crossval] Building {label_b} ({file_b.name})...")
    exit_b, output_b, hash_b = build_independent(file_b, target)
    build_b_ok = exit_b == 0

    return {
        "models": {
            label_a: {"file": str(file_a), "lean_build_hash": hash_a, "build_ok": build_a_ok},
            label_b: {"file": str(file_b), "lean_build_hash": hash_b, "build_ok": build_b_ok},
        },
        "summary": {
            "theorems_a": len(theorems_a),
            "theorems_b": len(theorems_b),
            "common": len(common),
            "unique_to_a": list(only_a),
            "unique_to_b": list(only_b),
            "equivalent": sum(1 for r in equiv_results if r["equivalent"]),
            "non_equivalent": sum(1 for r in equiv_results if not r["equivalent"]),
        },
        "theorems": equiv_results,
        "all_equivalent": all_equivalent,
        "build_result": "PASS" if build_a_ok and build_b_ok else "FAIL",
        "verdict": "PASS" if all_equivalent and build_a_ok and build_b_ok else "FAIL",
    }


def main():
    parser = argparse.ArgumentParser(description="Layer 1: Multi-Model Cross-Validation")
    parser.add_argument("--model-a", type=Path, required=True, help="First Lean proof file")
    parser.add_argument("--model-b", type=Path, required=True, help="Second Lean proof file")
    parser.add_argument("--label-a", default="model_a", help="Label for first model")
    parser.add_argument("--label-b", default="model_b", help="Label for second model")
    parser.add_argument("--target", default="SilverSightRRC", help="lake build target")
    parser.add_argument("--timeout", type=int, default=120, help="Build timeout (seconds)")
    parser.add_argument("--output", type=Path, default=None, help="Output receipt path")
    args = parser.parse_args()

    result = cross_validate(
        args.model_a, args.model_b, args.target,
        args.label_a, args.label_b, args.timeout,
    )

    # Print summary
    s = result["summary"]
    print(f"\n  {s['theorems_a']} theorems in A, {s['theorems_b']} in B")
    print(f"  {s['common']} common, {len(s['unique_to_a'])} unique to A, {len(s['unique_to_b'])} unique to B")
    print(f"  {s['equivalent']} equivalent, {s['non_equivalent']} non-equivalent")
    print(f"  Build: {result['build_result']}")

    if not result["all_equivalent"]:
        print("\n  Non-equivalent theorems:")
        for t in result["theorems"]:
            if not t["equivalent"]:
                print(f"    ❌ {t['name']}")
                print(f"       A: {t['sig_a'][:80]}")
                print(f"       B: {t['sig_b'][:80]}")
        sys.exit(1)

    print(f"  ✅ All common theorems equivalent")
    print(f"  ✅ Both builds pass")
    print(f"  Verdict: {result['verdict']}")

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(result, indent=2))
        print(f"  Receipt: {args.output}")

    sys.exit(0 if result["verdict"] == "PASS" else 1)


if __name__ == "__main__":
    main()
