#!/usr/bin/env python3
"""Worker process for MCP autoproof. Handles individual proof requests."""

import os
import sys
import json
import subprocess
import time
import fcntl
import urllib.request
from pathlib import Path

SILVERSIGHT = Path(__file__).resolve().parent.parent
NEON_API = "http://100.92.88.64:8766/generate"
LOCK_DIR = SILVERSIGHT / ".lake" / "autoproof_worker_locks"
LOCK_DIR.mkdir(parents=True, exist_ok=True)

def acquire_lock(name: str, timeout: float = 5.0) -> bool:
    """Acquire a lock file."""
    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

def call_phi4(prompt: str) -> str:
    """Call LLM on neon-64gb."""
    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 response."""
    import re
    m = re.search(r'```lean\n(.*?)```', text, re.DOTALL)
    if m: return m.group(1).strip()
    lines = text.split('\n')
    return '\n'.join(l for l in lines if l and not l.startswith('```'))

def handle_fill_sorry(filepath: str) -> dict:
    """Handle fill_sorry request in worker process."""
    full = SILVERSIGHT / filepath
    if not full.exists():
        return {"status": "error", "message": f"File not found: {full}"}
    
    # Find sorry and get context
    content = full.read_text().splitlines(keepends=True)
    lines = [l.rstrip() for l in content]
    
    for i, line in enumerate(lines):
        if "sorry" in line and not line.strip().startswith("--"):
            start = max(0, i - 10)
            end = min(len(lines), i + 11)
            context = "\n".join(lines[start:end])
            break
    else:
        return {"status": "ok", "message": "No sorries found"}
    
    # Call LLM
    prompt = f"Complete this Lean 4 proof:\n{context}\n\nProof:"
    try:
        response = call_phi4(prompt)
        proof = extract_lean(response)
        if not proof:
            return {"status": "error", "message": "Empty proof from LLM"}
        return {"status": "success", "proof": proof[:200], "context": context[:50]}
    except Exception as e:
        return {"status": "error", "message": f"LLM failed: {e}"}

def main():
    """Main worker loop."""
    worker_id = None
    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == "--worker-id":
            worker_id = sys.argv[i+1]
        elif sys.argv[i] == "--stdio":
            # MCP stdio mode
            lock_name = f"worker:{worker_id or 'main'}"
            if not acquire_lock(lock_name):
                print(json.dumps({"status": "error", "message": "Worker locked"}))
                return
            
            try:
                for line in sys.stdin:
                    line = line.strip()
                    if not line: continue
                    try:
                        req = json.loads(line)
                        method = req.get("method", "")
                        params = req.get("params", {}).get("arguments", {})
                        
                        if method == "tools/call" and params.get("name") == "fill_sorry":
                            result = handle_fill_sorry(params.get("file", ""))
                            resp = {"id": req.get("id"), "result": {"content": [{"text": json.dumps(result)}]}}
                        else:
                            resp = {"id": req.get("id"), "error": {"code": -32601, "message": "Method not found"}}
                        
                        sys.stdout.write(json.dumps(resp) + "\n")
                        sys.stdout.flush()
                    except: pass
            finally:
                release_lock(lock_name)
        i += 1

if __name__ == "__main__":
    main()