#!/usr/bin/env python3
"""
MCP server for automated Lean proof filling.
Provides thread-safe operations via file-based locking.

Tools:
  - fill_sorry: Fill a sorry in a Lean file, build, keep or discard
  - check_proof: Run lake build and return errors
  - get_sorry_context: Get theorem context around a sorry
"""
import os
import sys
import json
import subprocess
import time
import fcntl
from pathlib import Path
from threading import Lock
from datetime import datetime

SILVERSIGHT = Path(__file__).resolve().parent.parent
NEON_API = "http://100.92.88.64:8766/generate"

# ── Locking ─────────────────────────────────────────────────────────────────

LOCK_DIR = SILVERSIGHT / ".lake" / "autoproof_locks"
LOCK_DIR.mkdir(parents=True, exist_ok=True)

def acquire_lock(name: str, timeout: float = 5.0) -> bool:
    """Acquire a lock file, return True if successful."""
    lockpath = LOCK_DIR / f"{name}.lock"
    try:
        lockpath.touch(exist_ok=False)
    except FileExistsError:
        pass
    
    fd = os.open(str(lockpath), os.O_RDWR | os.O_CREAT)
    start = time.time()
    while time.time() - start < timeout:
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            return True
        except (IOError, OSError):
            time.sleep(0.1)
    os.close(fd)
    return False

def release_lock(name: str):
    """Release a lock file."""
    lockpath = LOCK_DIR / f"{name}.lock"
    try:
        fd = os.open(str(lockpath), os.O_RDWR)
        fcntl.flock(fd, fcntl.LOCK_UN)
        os.close(fd)
    except: pass

# ── Worker Pool ─────────────────────────────────────────────────────────────

WORKER_DIR = SILVERSIGHT / ".lake" / "autoproof_workers"
WORKER_DIR.mkdir(parents=True, exist_ok=True)
MAX_WORKERS = 4
ACTIVE_WORKERS = []

def spawn_worker(worker_id: int):
    """Spawn a background worker process for handling proofs."""
    cmd = [
        sys.executable,
        str(SILVERSIGHT / "scripts" / "mcp_worker.py"),
        "--worker-id", str(worker_id),
        "--stdio"
    ]
    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=SILVERSIGHT
    )
    ACTIVE_WORKERS.append(proc)
    return proc

def get_worker_count() -> int:
    """Count active workers."""
    return len([p for p in ACTIVE_WORKERS if p.poll() is None])

def maybe_spawn_workers():
    """Spawn more workers if under load (queue depth > available workers)."""
    queue_depth = len(ACTIVE_WORKERS)
    if queue_depth < MAX_WORKERS:
        # Check if we have high concurrency demand
        if queue_depth > 1:
            for i in range(MAX_WORKERS - queue_depth):
                spawn_worker(queue_depth + i)
                time.sleep(0.1)

# ── Lean Utils ────────────────────────────────────────────────────────────

def find_sorry(filepath: str):
    """Find the first sorry in a file and return context."""
    full = SILVERSIGHT / filepath
    if not full.exists():
        return None, None, None
    
    content = full.read_text().splitlines(keepends=True)
    lines = [l.rstrip() for l in content]
    
    # Find sorry positions
    for i, line in enumerate(lines):
        if "sorry" in line and not line.strip().startswith("--"):
            # Extract context (theorem + 10 lines before, sorry + 10 lines after)
            start = max(0, i - 10)
            end = min(len(lines), i + 11)
            return content, "\n".join(lines[start:end]), content
    return content, None, content

def call_phi4(prompt: str) -> str:
    """Call the LLM on neon-64gb."""
    import urllib.request
    import urllib.parse
    
    data = json.dumps({"message": prompt, "max_tokens": 2000}).encode()
    req = urllib.request.Request(
        NEON_API,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read().decode())["response"]

def extract_lean(text: str) -> str:
    """Extract Lean code from LLM response."""
    import re
    # Try to find code blocks
    m = re.search(r'```lean\n(.*?)\n```', text, re.DOTALL)
    if m: return m.group(1).strip()
    # Try to find just proof text
    lines = text.split('\n')
    return '\n'.join(l for l in lines if l and not l.startswith('```'))

# ── MCP Tools ───────────────────────────────────────────────────────────────

def call_fill_sorry(filepath: str):
    """Fill a sorry with concurrency protection."""
    full = SILVERSIGHT / filepath
    if not full.exists():
        return {"status": "error", "message": f"File not found: {full}"}
    
    # Acquire file lock
    if not acquire_lock(f"file:{filepath}"):
        return {"status": "error", "message": f"File {filepath} is locked"}
    
    try:
        content, context, original = find_sorry(filepath)
        if context is None:
            return {"status": "ok", "message": "No sorries found", "sorries": 0}
        
        prompt = (
            "Complete this Lean 4 proof. Output ONLY the proof body.\n\n"
            f"{context}\n\nProof:"
        )
        
        try:
            response = call_phi4(prompt)
        except Exception as e:
            return {"status": "error", "message": f"LLM failed: {e}"}
        
        proof = extract_lean(response)
        if not proof:
            return {"status": "error", "message": "Empty proof from LLM"}
        
        # Apply proof to file
        new_content = content.replace("  sorry", f"  {proof}", 1)
        full.write_text(new_content)
        
        # Acquire build lock and run lake build
        if not acquire_lock("build", timeout=10.0):
            full.write_text(original)
            return {"status": "error", "message": "Build system locked"}
        
        try:
            result = subprocess.run(
                ["lake", "build", "CoreFormalism.CRTSidon"],
                cwd=SILVERSIGHT, capture_output=True, text=True, timeout=300
            )
            success = result.returncode == 0
        finally:
            release_lock("build")
        
        if success:
            # Commit changes
            if acquire_lock("git", timeout=10.0):
                try:
                    subprocess.run(["git", "add", "-f", str(filepath)],
                                  cwd=SILVERSIGHT, capture_output=True)
                    subprocess.run(["git", "commit", "-m",
                                    f"autoproof(mcp): {datetime.now().isoformat()}:{filepath}"],
                                  cwd=SILVERSIGHT, capture_output=True)
                finally:
                    release_lock("git")
            return {"status": "success", "proof": proof[:200], "errors": 0}
        else:
            full.write_text(original)
            return {"status": "failed", "proof": proof[:200],
                    "errors": result.returncode,
                    "output": result.stdout[-500:] + result.stderr[-500:]}
    finally:
        release_lock(f"file:{filepath}")

def call_check_proof(target: str):
    """Check proof with concurrency protection."""
    if not acquire_lock("build", timeout=10.0):
        return {"status": "error", "message": "Build system locked"}
    
    try:
        result = subprocess.run(
            ["lake", "build", target],
            cwd=SILVERSIGHT, capture_output=True, text=True, timeout=300
        )
        errors = result.stdout.count("error:") + result.stderr.count("error:")
        return {"status": "success" if result.returncode == 0 else "failed",
                "target": target, "errors": errors,
                "output": (result.stdout + result.stderr)[-500:]}
    finally:
        release_lock("build")

def call_get_sorry_context(filepath: str):
    """Get sorry context."""
    _, context, _ = find_sorry(filepath)
    if context is None:
        return {"status": "ok", "message": "No sorries found"}
    return {"status": "ok", "context": context}

# ── MCP Protocol ───────────────────────────────────────────────────────────

def handle_request(request):
    """Handle MCP request."""
    method = request.get("method", "")
    params = request.get("params", {})
    req_id = request.get("id", None)
    
    if method == "list_tools":
        return {"id": req_id, "result": {"tools": [
            {"name": "fill_sorry", "description": "Fill a sorry", "inputSchema": {"type": "object", "properties": {"file": {"type": "string"}}, "required": ["file"]}},
            {"name": "check_proof", "description": "Check proof", "inputSchema": {"type": "object", "properties": {"target": {"type": "string"}}, "required": ["target"]}},
            {"name": "get_sorry_context", "description": "Get context", "inputSchema": {"type": "object", "properties": {"file": {"type": "string"}}, "required": ["file"]}},
        ]}}
    
    elif method == "tools/call":
        tool = params.get("name", "")
        args = params.get("arguments", {})
        
        result = {"status": "error", "message": "Unknown tool"}
        if tool == "fill_sorry":
            result = call_fill_sorry(args.get("file", ""))
        elif tool == "check_proof":
            result = call_check_proof(args.get("target", ""))
        elif tool == "get_sorry_context":
            result = call_get_sorry_context(args.get("file", ""))
        
        return {"id": req_id, "result": {"content": [{"text": json.dumps(result)}]}}
    
    return {"id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}}

def main():
    """Main entry point."""
    if len(sys.argv) > 1 and sys.argv[1] == "--stdio":
        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue
            try:
                request = json.loads(line)
                response = handle_request(request)
                sys.stdout.write(json.dumps(response) + "\n")
                sys.stdout.flush()
            except: pass
    else:
        print("MCP server running. Use --stdio mode.")

if __name__ == "__main__":
    main()
