#!/usr/bin/env python3
"""
Gemma4-based Lean porting MCP server.
Finds TODO(lean-port) sorries, sends to Gemma for proof completion, validates.
"""
import json, os, re, subprocess, sys, urllib.request
from pathlib import Path
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("gemma-lean-port", log_level="WARNING")

GEMMA_URL = os.getenv("GEMMA_URL", "http://127.0.0.1:8081/v1/chat/completions")
GEMMA_MODEL = os.getenv("GEMMA_MODEL", "gemma4-12b")
LAKE_WORKDIR = os.getenv("LAKE_WORKDIR", "/home/allaun/SilverSight/0-Core-Formalism/lean/Semantics")


@mcp.tool()
def find_todo_sorries(repo_root: str = "") -> str:
    """Find all TODO(lean-port) sorry declarations in Lean files."""
    root = repo_root or LAKE_WORKDIR
    results = []
    for lean_file in Path(root).rglob("*.lean"):
        try:
            content = lean_file.read_text()
        except Exception:
            continue
        for m in re.finditer(r"TODO\(lean-port\)", content):
            line = content[:m.start()].count("\n") + 1
            results.append(f"{lean_file}:{line}")
    return "\n".join(results) or "No TODO(lean-port) found"


@mcp.tool()
def port_theorem(lean_file: str, line_no: int) -> str:
    """Send theorem context to Gemma and attempt proof. Returns result JSON."""
    lines = Path(lean_file).read_text().split("\n")
    start = max(0, line_no - 15)
    ctx = "\n".join(lines[start:line_no + 10])
    
    prompt = f"Complete this Lean theorem. Use Q16_16 fixed-point (no Float). Output ONLY the proof block starting with `:= by`.\n\n{ctx}"
    body = json.dumps({"model": GEMMA_MODEL, "messages": [{"role": "user", "content": prompt}]})
    req = urllib.request.Request(GEMMA_URL, data=body.encode(), headers={"Content-Type": "application/json"})
    
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            data = json.loads(resp.read())
            proof = data["choices"][0]["message"]["content"]
    except Exception as e:
        return json.dumps({"status": "error", "message": str(e)})
    
    # Extract proof and apply
    proof_text = extract_proof(proof)
    apply_proof(lean_file, line_no, proof_text)
    
    # Validate
    module = lean_file.split("/")[-1].replace(".lean", "")
    ok, err = validate_build(module)
    return json.dumps({"status": "success" if ok else "failed", "proof": proof_text[:200], "error": err[:200] if err else None})


def extract_proof(text: str) -> str:
    for line in text.split("\n"):
        s = line.strip()
        if ":= by" in s:
            return s.split(":= by", 1)[-1].strip()
    return text.strip()


def apply_proof(file: str, line_no: int, proof: str) -> bool:
    lines = Path(file).read_text().split("\n")
    for i, l in enumerate(lines):
        if "sorry" in l and abs(i - line_no + 1) <= 2:
            indent = len(l) - len(l.lstrip())
            lines[i] = l.replace("sorry", ":=\n" + "\n".join(" " * (indent + 2) + p for p in proof.split("\n")))
            break
    Path(file).write_text("\n".join(lines))
    return True


def validate_build(module: str) -> tuple[bool, str]:
    result = subprocess.run(["lake", "build", module], cwd=LAKE_WORKDIR, capture_output=True, text=True, timeout=180)
    ok = result.returncode == 0 and "sorry" not in (result.stdout + result.stderr)
    return ok, (result.stdout + result.stderr)[:500]


if __name__ == "__main__":
    mcp.run(transport="stdio")