#!/usr/bin/env python3
"""
prove.py — Lean proof-filling loop with 5-check verification pipeline.

Adapted from Peng et al. (2026) pipeline-math verify.sh:
  1. SHA pin check    — frozen theorem stubs haven't drifted
  2. Banned keywords  — no sorry/native_decide/admit/axiom in proof files
  3. lake build clean — 0 errors, 0 unexpected warnings
  4. #print axioms    — proof depends only on {propext, Classical.choice, Quot.sound}
  5. Discharge gate   — @Frozen = @Proof := rfl (proof matches theorem type)

Metric: build_errors * 1000 + sorries (lower = better, 0 = perfect).
"""
import os, sys, time, json, subprocess, re, hashlib, shutil
from pathlib import Path
import urllib.request

# ── Config ──────────────────────────────────────────────────────────────

SILVERSIGHT_DIR = os.path.expanduser("~/SilverSight")
LEAN_FILE = os.environ.get("TARGET_FILE", "formal/CoreFormalism/Bind.lean")
LLM_MODEL = os.environ.get("LLM_MODEL", "phi4:14b")
LLM_BASE_URL = os.environ.get("LLM_BASE_URL", "http://100.92.88.64:11434/v1")
LLM_API_KEY = os.environ.get("LLM_API_KEY", "sk-local")
TIME_BUDGET_S = int(os.environ.get("TIME_BUDGET_S", "600"))

BANNED_KEYWORDS = ["sorryAx", "native_decide", "admit", "unsafe",
                   "implemented_by", "ofReduceBool"]
ALLOWED_AXIOMS = {"propext", "Classical.choice", "Quot.sound"}

# ── LLM client ─────────────────────────────────────────────────────────

def llm_complete(prompt, system="You are a Lean 4 proof engineer."):
    body = json.dumps({
        "model": LLM_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.3,
        "max_tokens": 4096,
    }).encode()
    req = urllib.request.Request(
        f"{LLM_BASE_URL}/chat/completions",
        data=body,
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {LLM_API_KEY}"},
        method="POST",
    )
    try:
        resp = urllib.request.urlopen(req, timeout=120)
        data = json.loads(resp.read())
        return data["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"LLM error: {e}")
        return None

# ─── Check 1: SHA pins ────────────────────────────────────────────────

def check_sha_pins(lean_file):
    """Check SHA-256 of file against pinned hash in scripts/frozen.sha256."""
    pins_path = Path(SILVERSIGHT_DIR) / "scripts" / "frozen.sha256"
    if not pins_path.exists():
        return True, "no pins file (SKIP)"
    rel_path = os.path.relpath(
        os.path.join(SILVERSIGHT_DIR, lean_file),
        SILVERSIGHT_DIR
    )
    content = pins_path.read_text()
    for line in content.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split(None, 1)
        if len(parts) == 2 and parts[1] == rel_path:
            pinned = parts[0]
            actual = sha256_of(os.path.join(SILVERSIGHT_DIR, lean_file))
            if pinned == actual:
                return True, f"SHA pin match: {rel_path}"
            else:
                return False, f"SHA pin MISMATCH {rel_path}: pinned={pinned[:16]} actual={actual[:16]}"
    return True, "no pin for this file (SKIP)"

def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        h.update(f.read())
    return h.hexdigest()

# ─── Check 2: Banned keywords ─────────────────────────────────────────

def check_banned_keywords(text, allow_sorry=False):
    """Check for banned keywords in Lean code. Returns (pass, message)."""
    cleaned = strip_comments(text)
    for kw in BANNED_KEYWORDS:
        if re.search(r'\b' + re.escape(kw) + r'\b', cleaned):
            return False, f"banned keyword: {kw}"
    if not allow_sorry:
        if re.search(r'\bsorry\b', cleaned):
            return False, "banned keyword: sorry (not allowed in proof files)"
    return True, "no banned keywords"

def strip_comments(s):
    """Strip Lean comments (-- line, /- ... -/ block)."""
    out = []
    i, n = 0, len(s)
    depth = 0
    while i < n:
        two = s[i:i+2]
        if depth == 0 and two == "--":
            j = s.find("\n", i)
            i = j if j != -1 else n
        elif two == "/-":
            depth += 1
            i += 2
        elif depth > 0 and two == "-/":
            depth -= 1
            i += 2
        elif depth > 0:
            i += 1
        else:
            out.append(s[i])
            i += 1
    return "".join(out)

# ─── Check 3: lake build ──────────────────────────────────────────────

def lake_build(target=None, timeout=120):
    """Run `lake build` and return (error_count, output)."""
    cmd = ["lake", "build"]
    if target:
        cmd.append(target)
    try:
        result = subprocess.run(
            cmd, cwd=SILVERSIGHT_DIR, capture_output=True, text=True,
            timeout=timeout,
        )
        output = result.stdout + result.stderr
        errors = output.count("error:")
        warnings = output.count("warning:") - output.count("declaration uses .sorry.")
        return errors, warnings, output
    except subprocess.TimeoutExpired:
        return 999, 0, "TIMEOUT"
    except Exception as e:
        return 999, 0, str(e)

# ─── Check 4: #print axioms ───────────────────────────────────────────

def check_axioms(module_name):
    """Check that a module's axioms are within the allowed set."""
    ax_file = os.path.join(SILVERSIGHT_DIR, ".lake", "tmp_ax_check.lean")
    os.makedirs(os.path.dirname(ax_file), exist_ok=True)
    with open(ax_file, "w") as f:
        f.write(f"import {module_name}\n")
        f.write(f"#print axioms {module_name}\n")

    try:
        result = subprocess.run(
            ["lake", "env", "lean", ax_file],
            cwd=SILVERSIGHT_DIR, capture_output=True, text=True, timeout=60,
        )
        output = result.stdout + result.stderr
        os.remove(ax_file)
        # Parse axiom line
        m = re.search(r"'([^']+)' depends on axioms:\s*\[([^\]]+)\]", output)
        if m:
            axioms = set(a.strip() for a in m.group(2).split(",") if a.strip())
            bad = axioms - ALLOWED_AXIOMS
            if bad:
                return False, f"unexpected axioms: {bad}"
            return True, f"axioms OK: {axioms}"
        return True, "no axiom line (module might be empty)"
    except Exception as e:
        return False, f"axiom check error: {e}"

# ─── Check 5: Discharge gate ──────────────────────────────────────────

def check_discharge_gate(lean_file):
    """Check that Discharge.lean compiles — proof ↔ theorem type gate."""
    discharge_path = os.path.join(
        os.path.dirname(os.path.join(SILVERSIGHT_DIR, lean_file)),
        "Discharge.lean"
    )
    if not os.path.exists(discharge_path):
        return True, "no Discharge.lean (SKIP)"
    discharge_module = "Discharge"
    err, warn, out = lake_build(target=discharge_module, timeout=60)
    if err == 0:
        return True, "Discharge gate PASS"
    return False, f"Discharge gate FAIL ({err} errors)"

# ── Proof filling ──────────────────────────────────────────────────────

def count_sorries(text):
    cleaned = strip_comments(text)
    return len(re.findall(r'\bsorry\b', cleaned))

def get_sorry_context(filepath):
    path = Path(os.path.join(SILVERSIGHT_DIR, filepath))
    if not path.exists():
        return None, None, 0, "file not found"
    text = path.read_text()
    n = count_sorries(text)
    if n == 0:
        return text, None, 0, "no sorries found"
    # Find the first theorem/lemma that contains a sorry
    blocks = list(re.finditer(
        r'(theorem\s+\w+|lemma\s+\w+|def\s+\w+)[\s\S]*?(?=\b(theorem|lemma|def)\b|\Z)',
        text
    ))
    for m in blocks:
        block = m.group(0)
        if 'sorry' in block:
            return text, block, n, None
    return text, None, n, "could not find sorry block"

def fill_sorry(filepath, proof_text):
    path = Path(os.path.join(SILVERSIGHT_DIR, filepath))
    text = path.read_text()
    new_text = text.replace("sorry", proof_text.strip(), 1)
    path.write_text(new_text)
    return new_text

def restore_file(filepath, original):
    path = Path(os.path.join(SILVERSIGHT_DIR, filepath))
    path.write_text(original)

# ── Main experiment loop ────────────────────────────────────────────────

def experiment():
    print("=" * 60)
    print("  AutoResearch: Lean Proof Filling (5-check pipeline)")
    print(f"  Model: {LLM_MODEL}")
    print(f"  File: {LEAN_FILE}")
    print(f"  Budget: {TIME_BUDGET_S}s")
    print("=" * 60)

    t_start = time.time()
    t_end = t_start + TIME_BUDGET_S
    results = []
    best_score = 999999

    initial_sorries = count_sorries(Path(os.path.join(SILVERSIGHT_DIR, LEAN_FILE)).read_text())
    print(f"  Initial sorries: {initial_sorries}")
    print()

    while time.time() < t_end:
        iteration = len(results) + 1
        t_iter = time.time()

        text, context, n, err = get_sorry_context(LEAN_FILE)
        if err:
            current = count_sorries(Path(os.path.join(SILVERSIGHT_DIR, LEAN_FILE)).read_text())
            if current == 0:
                print(f"\n*** All {initial_sorries} sorries filled! ***")
                break
            print(f"\n[{iteration}] {err}")
            break

        # Ask LLM to fill the sorry
        sys_prompt = "You are a Lean 4 proof engineer. Write complete, compilable proofs."
        prompt = f"""Fill the `sorry` in this Lean 4 theorem.
Write ONLY the proof body — the text that replaces `sorry` after `:= by`.
Do NOT repeat the theorem header. Use `by` block syntax, not `begin`/`end`.
Output raw Lean code only, no markdown.

{context}

Proof body to replace `sorry`:"""

        proof = llm_complete(prompt, sys_prompt)
        if not proof:
            print(f"  [{iteration}] LLM returned nothing, skipping")
            results.append({"iteration": iteration, "error": "no response"})
            continue

        # Post-process LLM output
        proof = proof.strip()
        if proof.startswith("```"):
            proof = proof.split("```")[1] if "```" in proof[3:] else proof
            if proof.startswith("lean"):
                proof = proof[4:]
        proof = proof.strip().rstrip("`").strip()
        if proof.startswith("theorem") or proof.startswith("lemma") or proof.startswith("def"):
            lines = proof.split("\n")
            for i, line in enumerate(lines):
                if line.strip().startswith(":= by"):
                    proof = "\n".join(lines[i+1:])
                    break
            else:
                proof = lines[-1]

        # Apply proof
        original = Path(os.path.join(SILVERSIGHT_DIR, LEAN_FILE)).read_text()
        fill_sorry(LEAN_FILE, proof)

        # ── 5-check pipeline ────────────────────────────────────────
        checks_passed = True
        new_text = Path(os.path.join(SILVERSIGHT_DIR, LEAN_FILE)).read_text()

        # Check 1: SHA pins
        pin_ok, pin_msg = check_sha_pins(LEAN_FILE)
        if not pin_ok:
            print(f"  [{iteration}] CHECK 1 FAIL: {pin_msg}")
            checks_passed = False

        # Check 2: Banned keywords
        kw_ok, kw_msg = check_banned_keywords(new_text, allow_sorry=False)
        if not kw_ok:
            print(f"  [{iteration}] CHECK 2 FAIL: {kw_msg}")
            checks_passed = False

        # Check 3: lake build
        build_errors, build_warnings, build_out = lake_build(timeout=120)
        current_sorries = count_sorries(new_text)
        elapsed = time.time() - t_iter

        score = build_errors * 1000 + current_sorries
        status = "KEEP" if checks_passed and score < best_score else "DISCARD"

        if status == "KEEP":
            # Check 4: #print axioms (only if build passed)
            if build_errors == 0:
                module_name = LEAN_FILE.replace("/", ".").replace(".lean", "")
                ax_ok, ax_msg = check_axioms(module_name)
                if not ax_ok:
                    print(f"  [{iteration}] CHECK 4 FAIL: {ax_msg}")
                    checks_passed = False
                    status = "DISCARD"

            # Check 5: Discharge gate
            if checks_passed:
                dg_ok, dg_msg = check_discharge_gate(LEAN_FILE)
                if not dg_ok:
                    print(f"  [{iteration}] CHECK 5 FAIL: {dg_msg}")
                    checks_passed = False
                    status = "DISCARD"

        if status == "KEEP":
            best_score = score
            subprocess.run(
                ["git", "add", "-f", LEAN_FILE],
                cwd=SILVERSIGHT_DIR, capture_output=True
            )
            subprocess.run(
                ["git", "commit", "-m",
                 f"autoresearch: {build_errors} err/{current_sorries} sorries (iter {iteration})"],
                cwd=SILVERSIGHT_DIR, capture_output=True
            )

        if status == "DISCARD":
            restore_file(LEAN_FILE, original)

        checks_str = ""
        if not checks_passed:
            checks_str = " CHECKS FAILED"
        print(f"  [{iteration:3d}] build_err={build_errors} sorries={current_sorries:3d}"
              f" elapsed={elapsed:.1f}s [{status}]{checks_str}")

        results.append({
            "iteration": iteration,
            "sorries": current_sorries,
            "build_errors": build_errors,
            "status": status,
            "checks_passed": checks_passed,
            "elapsed_s": round(elapsed, 1),
        })

        if current_sorries == 0 and build_errors == 0:
            print(f"\n*** ALL PROOFS COMPLETE! ***")
            break

    total = time.time() - t_start
    print(f"\n{'='*60}")
    print(f"  Summary: {len(results)} iterations, {total:.0f}s")
    print(f"  Best score (errors*1000+sorries): {best_score}")
    print(f"  Model: {LLM_MODEL}")
    print(f"{'='*60}")

if __name__ == "__main__":
    experiment()