#!/usr/bin/env python3
"""
TokenSaver MCP — routes requests to free local compute first, falls through to paid API.

Free tier (no tokens):
  - verify_build:     qfox-1 build server (HTTP)
  - classify_proof:   rrc-watchdog (local Lean exe)
  - compile_check:    local lean --run via SSH build server
  - generate_proof:   neon's local Ollama models (CPU, free)

Paid fallback (uses tokens):
  - generate_proof → OpenRouter DeepSeek V4 Flash (only if local models fail)

Usage:
  python3 token_saver_mcp.py
"""

from __future__ import annotations

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

from mcp.server.fastmcp import FastMCP

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

# ── Free resources ──────────────────────────────────────────────────────

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")
RRC_WATCHDOG = os.path.join(LAKE_WORKDIR, ".lake/build/bin/rrc-watchdog")

# ── Paid fallback ───────────────────────────────────────────────────────

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_MODEL = "deepseek/deepseek-v4-flash"

def _get_openrouter_key() -> str:
    auth_path = Path.home() / ".local" / "share" / "opencode" / "auth.json"
    try:
        auth = json.loads(auth_path.read_text())
        return auth.get("openrouter", {}).get("key", "")
    except Exception:
        return os.environ.get("OPENROUTER_API_KEY", "")

# ── Tools: Free tier ────────────────────────────────────────────────────

@mcp.tool()
def verify_build(module: str) -> str:
    """Verify a Lean module compiles. Uses local build server (free, 0 tokens)."""
    try:
        resp = urllib.request.urlopen(f"{BUILD_SERVER}/{module}", timeout=600)
        return resp.read().decode()
    except Exception as e:
        return json.dumps({"passed": False, "error": str(e)[:200]})


@mcp.tool()
def classify_proof(pist_label: str = "", rrc_shape: str = "logogramProjection") -> str:
    """Classify a proof through the RRC alignment gate. Local exe (free, 0 tokens)."""
    if not os.path.exists(RRC_WATCHDOG):
        return json.dumps({"error": "rrc-watchdog not built", "score": 0})
    try:
        result = subprocess.run(
            [RRC_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:
    """Compile-check a Lean snippet. Uses build server (free, 0 tokens).
    Writes code to a temp file and runs `lake build` on it.
    """
    import tempfile
    tmp = Path(tempfile.mkdtemp()) / "check.lean"
    tmp.write_text(code)
    try:
        result = subprocess.run(
            ["lake", "env", "lean", str(tmp)],
            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)[:1000]
        return json.dumps({"passed": ok, "errors": errors})
    except Exception as e:
        return json.dumps({"passed": False, "errors": str(e)[:200]})
    finally:
        tmp.unlink(missing_ok=True)


@mcp.tool()
def generate_proof(theorem_context: str, line_no: int = 0) -> str:
    """Generate a Lean proof. Tries free local models first, falls back to paid API.

    Free tier: neon's DeepSeek-Prover-V2 (CPU, 0 token cost)
    Paid fallback: OpenRouter DeepSeek V4 Flash (uses API tokens)
    """
    # Try free: neon Ollama (raw completion mode)
    proof = _try_neon(theorem_context)
    if proof and not proof.startswith("#"):
        return proof

    # Fallback: paid OpenRouter
    return _try_openrouter(theorem_context, line_no)


# ── Internals ───────────────────────────────────────────────────────────

def _try_neon(theorem_context: str) -> str:
    """Try generating a proof using neon's free local model."""
    prompt = f"theorem {theorem_context} := by"
    body = json.dumps({
        "model": NEON_MODEL,
        "prompt": prompt,
        "raw": True,
        "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:
            data = json.loads(resp.read())
        text = data.get("response", "")
        return _extract_proof(text)
    except Exception as e:
        return f"# neon error: {e}"


def _try_openrouter(theorem_context: str, line_no: int) -> str:
    """Fallback: generate proof via OpenRouter (token cost)."""
    key = _get_openrouter_key()
    if not key:
        return "# No OpenRouter key configured"

    prompt = (
        "You are a Lean 4 theorem prover. "
        "Output ONLY the proof block starting with `:= by`. "
        "No markdown fences. No explanation.\n\n"
        f"Complete this Lean theorem:\n\n{theorem_context}"
    )

    body = json.dumps({
        "model": OPENROUTER_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.3, "max_tokens": 4096,
    }).encode()

    req = urllib.request.Request(
        OPENROUTER_URL, data=body,
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {key}"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read())
        content = data["choices"][0]["message"]["content"]
        return _extract_proof(content) or content[:500]
    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:
            if s.startswith("```"):
                continue
            proof_lines.append(s)
    if proof_lines:
        return "\n".join(proof_lines)
    for line in lines:
        s = line.strip()
        if s.startswith(":= by") or s.startswith("by "):
            return s
    return ""


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