#!/usr/bin/env python3
"""
TokenSaver MCP — intercepts token-burning operations, routes to free local compute.

Free tier (runs on neon CPU or qfox-1 build server, 0 token cost):
  - verify_build      lake build verification
  - classify_proof    RRC alignment classification
  - compile_check     Lean snippet compilation
  - linter_check      Lean linter pass
  - generate_proof    DeepSeek-Prover on neon (CPU, free)
  - format_code       Lean code formatting (lake fmt)
  - doc_lookup        Mathlib doc search (local cache)

Paid fallback (only if free tier fails):
  - generate_proof → OpenRouter DeepSeek V4 Flash

Usage: python3 token_saver_mcp.py
Register in opencode.json as a local MCP server.
Any MCP-compatible agent (Claude Code, OpenCode, Cursor) can call it.
"""

from __future__ import annotations

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

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("token-saver", log_level="WARNING")


def _ensure_singleton() -> None:
    """Prevent multiple instances of this MCP server from running simultaneously."""
    lock_path = Path.home() / ".local" / "share" / "opencode" / "token_saver_mcp.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR)
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except (OSError, BlockingIOError):
        print("token_saver_mcp is already running; exiting singleton instance.", file=sys.stderr)
        sys.exit(0)


# ── Resource URLs (configurable via env) ────────────────────────────────

BUILD_SERVER = os.environ.get("BUILD_SERVER_URL", "http://100.88.57.96:8765")
NEON_OLLAMA   = os.environ.get("NEON_OLLAMA_URL", "http://100.92.88.64:11434")
NEON_MODEL    = os.environ.get("NEON_MODEL",
    "hf.co/irmma/DeepSeek-Prover-V2-7B-Q4_K_M-GGUF")
LAKE_WORKDIR  = os.environ.get("LAKE_WORKDIR",
    "/home/allaun/Research Stack/0-Core-Formalism/lean/Semantics")

# ── Registered tools ────────────────────────────────────────────────────

TOOL_MANIFEST = """
All tools in this server run on free local compute (neon CPU / qfox-1 build server).
Zero API tokens consumed unless explicitly marked [PAID].

  verify_build(module)        → build server (free)
  classify_proof(label,shape) → rrc-watchdog (free)
  compile_check(code)         → lake env lean (free)
  linter_check(code)          → lake build --warnings (free)
  generate_proof(context)     → neon Ollama → OpenRouter fallback [free → PAID]
  format_code(code)           → lake fmt (free)
  doc_lookup(query)           → local mathlib doc cache (free)
"""


def _call_build_server(module: str, timeout: int = 600) -> str:
    """Call the local build server. 0 tokens."""
    try:
        resp = urllib.request.urlopen(f"{BUILD_SERVER}/{module}", timeout=timeout)
        return resp.read().decode()
    except Exception as e:
        return json.dumps({"passed": False, "error": str(e)[:200]})


def _call_neon(prompt: str, raw: bool = True) -> str:
    """Call neon's local Ollama model. 0 tokens."""
    body = json.dumps({
        "model": NEON_MODEL,
        "prompt": prompt,
        "raw": raw,
        "stream": False,
        "options": {"temperature": 0, "num_predict": 512, "num_thread": 18},
    }).encode()
    try:
        req = urllib.request.Request(
            f"{NEON_OLLAMA}/api/generate", data=body,
            headers={"Content-Type": "application/json"}, method="POST",
        )
        with urllib.request.urlopen(req, timeout=300) as resp:
            return json.loads(resp.read()).get("response", "")
    except Exception as e:
        return f"# neon error: {e}"


def _call_openrouter(prompt: str) -> str:
    """Fallback: paid OpenRouter API."""
    auth_path = Path.home() / ".local" / "share" / "opencode" / "auth.json"
    key = ""
    try:
        auth = json.loads(auth_path.read_text())
        key = auth.get("openrouter", {}).get("key", "")
    except Exception:
        key = os.environ.get("OPENROUTER_API_KEY", "")
    if not key:
        return "# No API key configured"

    body = json.dumps({
        "model": "deepseek/deepseek-v4-flash",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.3, "max_tokens": 4096,
    }).encode()
    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}

    try:
        req = urllib.request.Request(
            "https://openrouter.ai/api/v1/chat/completions", data=body,
            headers=headers, method="POST",
        )
        with urllib.request.urlopen(req, timeout=180) as resp:
            return json.loads(resp.read())["choices"][0]["message"]["content"]
    except Exception as e:
        return f"# OpenRouter error: {e}"


def _extract_proof(text: str) -> str:
    lines = text.split("\n")
    in_proof = False
    proof_lines: list[str] = []
    for line in lines:
        s = line.strip()
        if s.startswith(":= by"):
            in_proof = True
            proof_lines = [s]
            continue
        if in_proof and not s.startswith("```"):
            proof_lines.append(s)
    return "\n".join(proof_lines) if proof_lines else ""


# ═══════════════════════════════════════════════════════════════════════════
# Tools
# ═══════════════════════════════════════════════════════════════════════════

@mcp.tool()
def verify_build(module: str) -> str:
    """[FREE] Verify a Lean module compiles via local build server."""
    return _call_build_server(module)


@mcp.tool()
def classify_proof(pist_label: str = "", rrc_shape: str = "logogramProjection") -> str:
    """[FREE] Classify a proof through RRC alignment gate (local rrc-watchdog)."""
    watchdog = os.path.join(LAKE_WORKDIR, ".lake/build/bin/rrc-watchdog")
    if not os.path.exists(watchdog):
        return json.dumps({"error": "rrc-watchdog not built", "score": 0})
    try:
        result = subprocess.run(
            [watchdog, "--pist-label", pist_label or "none",
             "--exact-label", pist_label or "none", "--rrc-shape", rrc_shape],
            capture_output=True, text=True, timeout=30,
        )
        return result.stdout or json.dumps({"error": "empty", "score": 0})
    except Exception as e:
        return json.dumps({"error": str(e)[:200], "score": 0})


@mcp.tool()
def compile_check(code: str) -> str:
    """[FREE] Compile-check a Lean snippet locally. 0 tokens."""
    import tempfile, shutil
    tmp_dir = Path(tempfile.mkdtemp())
    try:
        tmp_file = tmp_dir / "check.lean"
        tmp_file.write_text(code)
        result = subprocess.run(
            ["lake", "env", "lean", str(tmp_file)],
            capture_output=True, text=True, timeout=120, cwd=LAKE_WORKDIR,
        )
        ok = result.returncode == 0
        errors = "\n".join(
            l for l in (result.stdout + result.stderr).split("\n")
            if "error:" in l or "warning:" in l
        )[:2000]
        return json.dumps({"passed": ok, "errors": errors})
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


@mcp.tool()
def linter_check(code: str) -> str:
    """[FREE] Run Lean linter on a snippet. 0 tokens."""
    import tempfile, shutil
    tmp_dir = Path(tempfile.mkdtemp())
    try:
        tmp_file = tmp_dir / "check.lean"
        tmp_file.write_text(code)
        result = subprocess.run(
            ["lake", "env", "lean", f"--lint={str(tmp_file)}"],
            capture_output=True, text=True, timeout=120, cwd=LAKE_WORKDIR,
        )
        return json.dumps({
            "passed": result.returncode == 0,
            "warnings": "\n".join(
                l for l in (result.stdout + result.stderr).split("\n")
                if "warning:" in l
            )[:2000],
        })
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


@mcp.tool()
def generate_proof(theorem_context: str, line_no: int = 0) -> str:
    """[FREE → PAID] Generate a Lean proof. Tries neon's CPU model first (0 tokens),
    falls back to OpenRouter DeepSeek V4 Flash (token cost)."""
    # Free tier: neon's DeepSeek-Prover with raw completion
    prompt = f"theorem {theorem_context} := by"
    result = _call_neon(prompt)
    proof = _extract_proof(result)

    if proof:
        return proof

    # Paid fallback: OpenRouter
    paid_prompt = (
        "You are a Lean 4 theorem prover. "
        "Output ONLY the proof block starting with `:= by`. "
        "No markdown fences.\n\n"
        f"Complete:\n\n{theorem_context}"
    )
    result = _call_openrouter(paid_prompt)
    proof = _extract_proof(result)
    return proof or result[:500]


@mcp.tool()
def format_code(code: str) -> str:
    """[FREE] Format Lean code using lake fmt. 0 tokens."""
    import tempfile, shutil
    tmp_dir = Path(tempfile.mkdtemp())
    try:
        tmp_file = tmp_dir / "check.lean"
        tmp_file.write_text(code)
        result = subprocess.run(
            ["lake", "fmt", str(tmp_file)],
            capture_output=True, text=True, timeout=30, cwd=LAKE_WORKDIR,
        )
        if result.returncode == 0:
            return tmp_file.read_text()
        return code  # Return original if fmt fails
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


@mcp.tool()
def doc_lookup(query: str) -> str:
    """[FREE] Look up a Lean/mathlib symbol in local docs. 0 tokens.
    Uses a temp file because `lean` doesn't support inline -e evaluation."""
    import tempfile, shutil
    tmp_dir = Path(tempfile.mkdtemp())
    try:
        tmp_file = tmp_dir / "check.lean"
        tmp_file.write_text(f"#check {query}\n")
        result = subprocess.run(
            ["lake", "env", "lean", str(tmp_file)],
            capture_output=True, text=True, timeout=30, cwd=LAKE_WORKDIR,
        )
        output = (result.stdout + result.stderr)[:2000]
        if not output.strip():
            return f"Symbol '{query}' not found"
        # Strip absolute tmp path noise from error messages
        lines = []
        for line in output.split("\n"):
            cleaned = line.replace(str(tmp_dir) + "/", "")
            if cleaned.strip():
                lines.append(cleaned)
        return "\n".join(lines[-10:])  # last 10 lines = signal
    except Exception as e:
        return f"Error: {e}"
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


@mcp.tool()
def health() -> str:
    """Check which free resources are reachable."""
    results = {}
    # Build server
    try:
        urllib.request.urlopen(f"{BUILD_SERVER}/Semantics", timeout=5)
        results["build_server"] = "up"
    except Exception:
        results["build_server"] = "down"
    # Neon Ollama
    try:
        urllib.request.urlopen(f"{NEON_OLLAMA}/api/tags", timeout=5)
        results["neon_ollama"] = "up"
    except Exception:
        results["neon_ollama"] = "down"
    # rrc-watchdog
    watchdog = os.path.join(LAKE_WORKDIR, ".lake/build/bin/rrc-watchdog")
    results["rrc_watchdog"] = "present" if os.path.exists(watchdog) else "missing"
    return json.dumps(results)


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