#!/usr/bin/env python3
"""autoproof.py — Auto-fill `sorry` blocks in Lean using phi4 on neon-64gb.

Usage:
  python3 scripts/autoproof.py
  python3 scripts/autoproof.py --file formal/CoreFormalism/CRTSidon.lean
  python3 scripts/autoproof.py --all  # fill all sorries in the file
"""
import os
import sys
import re
import json
import subprocess
import urllib.request

NEON_API = "http://100.92.88.64:8766/generate"
SILVERSIGHT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def call_phi4(prompt, timeout=600):
    """Call phi4 on neon via the LeanCopilot-compatible server."""
    body = json.dumps({
        "name": "phi4:14b",
        "input": prompt,
    }).encode()
    req = urllib.request.Request(NEON_API, data=body,
        headers={"Content-Type": "application/json"})
    resp = urllib.request.urlopen(req, timeout=timeout)
    data = json.loads(resp.read())
    return data["outputs"][0]["text"]

def extract_lean(text):
    """Extract Lean code from LLM response (strip markdown fences)."""
    text = text.strip()
    if text.startswith("```"):
        # Remove first ``` line
        text = text.split("\n", 1)[1] if "\n" in text else text
        # Remove trailing ```
        text = text.rsplit("```", 1)[0].strip()
    # Remove leading "lean" or "lean4" language tag
    if text.startswith("lean"):
        text = text[4:].strip()
    # Remove theorem header if it was repeated
    lines = text.split("\n")
    cleaned = []
    in_proof = False
    for line in lines:
        if ":= by" in line:
            # Remove everything up to and including `:= by`
            line = line.split(":= by", 1)[1].strip()
            in_proof = True
        if in_proof:
            cleaned.append(line)
    if cleaned:
        text = "\n".join(cleaned)
    return text.strip()

def find_sorry_context(filepath):
    """Find the first `sorry` and return its context."""
    with open(filepath) as f:
        content = f.read()
    
    # Find the theorem/lemma containing the first sorry
    lines = content.split("\n")
    theorem_start = None
    sorry_line = None
    
    for i, line in enumerate(lines):
        if re.match(r'\s*(theorem|lemma)\s', line):
            theorem_start = i
        if "sorry" in line and theorem_start is not None:
            sorry_line = i
            break
    
    if sorry_line is None:
        return content, None, content  # no sorries
    
    # Extract the theorem block (from theorem to the sorry)
    context_lines = lines[theorem_start:sorry_line + 1]
    context = "\n".join(context_lines)
    return content, context, content

def replace_sorry(filepath, proof_text):
    """Replace the first `sorry` with proof_text."""
    with open(filepath) as f:
        content = f.read()
    
    # Replace the first occurrence of "  sorry" with the proof
    # We need proper indentation
    new_content = content.replace("  sorry", f"  {proof_text}", 1)
    
    with open(filepath, "w") as f:
        f.write(new_content)
    return new_content

def lake_build(target="CoreFormalism.CRTSidon"):
    """Run lake build and return (success, output)."""
    result = subprocess.run(
        ["lake", "build", target],
        cwd=SILVERSIGHT,
        capture_output=True, text=True,
        timeout=300,
    )
    output = result.stdout + result.stderr
    success = result.returncode == 0
    return success, output

def main():
    filepath = sys.argv[sys.argv.index("--file") + 1] if "--file" in sys.argv else \
        "formal/CoreFormalism/CRTSidon.lean"
    fill_all = "--all" in sys.argv
    
    full_path = os.path.join(SILVERSIGHT, filepath)
    print(f"=== AutoProof: {filepath} ===")
    
    content, context, _ = find_sorry_context(full_path)
    if context is None:
        print("No sorries found.")
        return
    
    sorry_count = content.count("sorry")
    print(f"Found {sorry_count} sorry/s.")
    
    while True:
        content, context, original = find_sorry_context(full_path)
        if context is None:
            print("All sorries filled!")
            break
        
        print(f"\n--- Filling sorry ---")
        print(f"Context:\n{context}\n")
        
        # Build prompt
        prompt = (
            "Fill the `sorry` in this Lean 4 theorem. "
            "Output ONLY the proof body that replaces `sorry` after `:= by`. "
            "No markdown, no theorem header, no commentary.\n\n"
            f"{context}\n\nProof body:"
        )
        
        # Call phi4 on neon
        print("Calling phi4 on neon-64gb...", flush=True)
        response = call_phi4(prompt)
        proof = extract_lean(response)
        
        if not proof:
            print("Empty proof from LLM, retrying...")
            continue
        
        print(f"Generated proof ({len(proof)} chars): {proof[:100]}...")
        
        # Save original and apply
        replace_sorry(full_path, proof)
        
        # Build
        print("Running lake build...", flush=True)
        success, output = lake_build()
        
        if success:
            print("  PASSED!")
            subprocess.run(["git", "add", "-f", filepath], cwd=SILVERSIGHT)
            subprocess.run(
                ["git", "commit", "-m", f"autoproof: filled sorry in {filepath}"],
                cwd=SILVERSIGHT,
            )
            print("  Committed.")
            if not fill_all:
                break
        else:
            print("  FAILED. Restoring original...")
            with open(full_path, "w") as f:
                f.write(original)
            print("  Restored.")
            # Could retry with different prompt
            break
    
    print("\n=== Done ===")

if __name__ == "__main__":
    main()