#!/usr/bin/env python3
"""
break_glass_mcp.py — Last-resort multi-model panel via OpenRouter Fusion.

Wraps OpenRouter Fusion as an MCP tool. Only invoked when all other
approaches have failed on a novel problem.

Usage:
  python3 break_glass_mcp.py   # starts MCP server on stdio

Requires: OPENROUTER_API_KEY env var
"""

import json
import os
import sys
from datetime import datetime

try:
    import httpx
except ImportError:
    httpx = None

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
OPENROUTER_BASE = "https://openrouter.ai/api/v1"

# Problem type → model panel configuration
PANELS = {
    "lean_proof": {
        "analysis_models": [
            "anthropic/claude-opus-4.8",      # Lean expertise, formal proof
            "deepseek/deepseek-v4-pro",        # Best math reasoning
            "openai/gpt-5.5",                  # Widest training data
            "cohere/command-a-03-2026",        # Structured output, tool-use optimized
            "google/gemini-3.1-pro-preview",   # Different thinking style
            "zhipu/glm-5.2",                   # 753B, genuinely different architecture
            "moonshotai/kimi-k2.6",            # Strong code, different training
            "deepseek/deepseek-v4-flash",      # Fast, good for quick iterations
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 2.50,
    },
    "code_bug": {
        "analysis_models": [
            "openai/gpt-5.5",                  # Breadth, ops knowledge
            "anthropic/claude-opus-4.8",       # Depth, precision
            "deepseek/deepseek-v4-pro",        # Math/code reasoning
            "moonshotai/kimi-k2.6",            # Different code style
            "cohere/command-a-03-2026",        # Structured output
            "deepseek/deepseek-v4-flash",      # Fast iterations
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 1.80,
    },
    "architecture": {
        "analysis_models": [
            "anthropic/claude-opus-4.8",       # Deep reasoning
            "openai/gpt-5.5",                  # Breadth
            "deepseek/deepseek-v4-pro",        # Math rigor
            "cohere/command-a-03-2026",        # Structured output
            "google/gemini-3.1-pro-preview",   # Different perspective
            "zhipu/glm-5.2",                   # Architectural diversity
            "moonshotai/kimi-k2.6",            # Code-first thinking
            "deepseek/deepseek-v4-flash",      # Quick iterations
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 2.50,
    },
    "math_novel": {
        "analysis_models": [
            "deepseek/deepseek-v4-pro",        # Best math reasoning
            "anthropic/claude-opus-4.8",       # Formal rigor
            "google/gemini-3.1-pro-preview",   # Breadth
            "cohere/command-a-03-2026",        # Structured output
            "zhipu/glm-5.2",                   # Different math tradition
            "deepseek/deepseek-v4-flash",      # Quick checks
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 2.00,
    },
    "infra_debug": {
        "analysis_models": [
            "openai/gpt-5.5",                  # Ops knowledge
            "anthropic/claude-opus-4.8",       # Precision
            "deepseek/deepseek-v4-pro",        # Systems reasoning
            "cohere/command-a-03-2026",        # Structured tool chains
            "deepseek/deepseek-v4-flash",      # Quick checks
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 1.50,
    },
    "hard_wall": {
        "analysis_models": [
            "anthropic/claude-opus-4.8",
            "deepseek/deepseek-v4-pro",
            "openai/gpt-5.5",
            "cohere/command-a-03-2026",
            "google/gemini-3.1-pro-preview",
            "zhipu/glm-5.2",
            "moonshotai/kimi-k2.6",
            "deepseek/deepseek-v4-flash",
        ],
        "judge": "anthropic/claude-opus-4.8",
        "budget": 3.50,
    },
}

SYSTEM_PROMPT = """You are a specialist agent assembled to attack a specific
novel problem that has resisted all常规 approaches. You have access to:

1. The full codebase context (provided below)
2. The project's operating constraints
3. The specific problem statement
4. Previous attempts and their failure modes

YOUR TASK:
- Diagnose the root cause
- Propose a concrete solution with exact file paths and code
- If the problem is undecidable, explain why and propose what would be needed

CONSTRAINTS:
- Lean is the source of truth for formal claims
- No Float in compute paths (use Q0_16 or Q16_16)
- No sorry in committed code unless with TODO(lean-port) + human sign-off
- All new formal work goes to SilverSight (not Research Stack)

OUTPUT FORMAT (JSON):
{
  "diagnosis": "...",
  "solution": {
    "approach": "...",
    "confidence": 0.0-1.0,
    "files_to_change": ["..."],
    "code_snippets": {"file": "code"}
  },
  "alternatives": ["..."],
  "risks": ["..."],
  "cost_estimate": "$X.XX"
}"""


def classify_problem(problem: str) -> str:
    """Auto-classify problem type from text."""
    lower = problem.lower()
    if any(w in lower for w in ["sorry", "lean", "theorem", "proof", "lake build"]):
        return "lean_proof"
    if any(w in lower for w in ["architecture", "design", "refactor", "structure"]):
        return "architecture"
    if any(w in lower for w in ["math", "formula", "eigenvalue", "convergence"]):
        return "math_novel"
    if any(w in lower for w in ["infra", "deploy", "k3s", "docker", "ssh"]):
        return "infra_debug"
    if any(w in lower for w in ["stuck", "wall", "impossible", "no idea", "tried everything"]):
        return "hard_wall"
    return "code_bug"


def call_fusion(problem: str, context: str, problem_type: str,
                budget: float = 1.00) -> dict:
    """Call OpenRouter Fusion with the assembled context."""
    if not OPENROUTER_API_KEY:
        return {"error": "OPENROUTER_API_KEY not set"}
    if httpx is None:
        return {"error": "httpx not installed (pip install httpx)"}

    panel = PANELS.get(problem_type, PANELS["code_bug"])

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"""## Problem Statement
{problem}

## Codebase Context
{context}

## Problem Type
{problem_type}

Please diagnose and propose a solution."""},
    ]

    payload = {
        "model": "openrouter/fusion",
        "messages": messages,
        "plugins": [{
            "id": "fusion",
            "model": panel["judge"],
            "analysis_models": panel["analysis_models"],
        }],
        "max_tokens": 8192,
    }

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://researchstack.info",
        "X-Title": "break-glass-agent",
    }

    try:
        with httpx.Client(timeout=120) as client:
            resp = client.post(
                f"{OPENROUTER_BASE}/chat/completions",
                json=payload,
                headers=headers,
            )
            resp.raise_for_status()
            data = resp.json()

        # Extract cost
        usage = data.get("usage", {})
        cost = usage.get("total_cost", 0)

        # Extract response
        content = data["choices"][0]["message"]["content"]

        return {
            "response": content,
            "cost": cost,
            "model": data.get("model", "openrouter/fusion"),
            "usage": usage,
        }
    except Exception as e:
        return {"error": str(e)}


def handle_tool_call(tool_name: str, arguments: dict) -> dict:
    """Handle an MCP tool call."""
    if tool_name != "break_glass":
        return {"error": f"Unknown tool: {tool_name}"}

    problem = arguments.get("problem", "")
    problem_type = arguments.get("problem_type") or classify_problem(problem)
    budget = arguments.get("budget", 1.00)
    context_files = arguments.get("context_files", [])
    context = arguments.get("context", "")

    # Cost gate
    panel = PANELS.get(problem_type, PANELS["code_bug"])
    if budget > panel["budget"] * 3:
        return {
            "error": f"Budget ${budget:.2f} exceeds threshold ${panel['budget'] * 3:.2f}. "
                     f"Escalate to human.",
            "problem_type": problem_type,
        }

    result = call_fusion(problem, context, problem_type, budget)

    return {
        "problem_type": problem_type,
        "panel": [panel["judge"]] + panel["analysis_models"],
        "result": result,
        "timestamp": datetime.utcnow().isoformat(),
    }


def mcp_server():
    """Simple MCP server on stdio (JSON-RPC)."""
    for line in sys.stdin:
        try:
            req = json.loads(line.strip())
        except json.JSONDecodeError:
            continue

        method = req.get("method", "")
        req_id = req.get("id")
        params = req.get("params", {})

        if method == "initialize":
            resp = {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {"tools": {}},
                    "serverInfo": {
                        "name": "break-glass",
                        "version": "0.1.0",
                    },
                },
            }
        elif method == "tools/list":
            resp = {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": {
                    "tools": [{
                        "name": "break_glass",
                        "description": (
                            "Last-resort multi-model panel for novel problems. "
                            "Costs 3-5x a single model call. "
                            "Requires explicit 'break glass:' trigger."
                        ),
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "problem": {
                                    "type": "string",
                                    "description": "The specific problem statement",
                                },
                                "context": {
                                    "type": "string",
                                    "description": "Relevant codebase context",
                                },
                                "problem_type": {
                                    "type": "string",
                                    "enum": list(PANELS.keys()),
                                    "description": "Problem classification (auto-detected if omitted)",
                                },
                                "budget": {
                                    "type": "number",
                                    "description": "Maximum cost in USD (default: 1.00)",
                                },
                            },
                            "required": ["problem"],
                        },
                    }],
                },
            }
        elif method == "tools/call":
            tool_name = params.get("name", "")
            arguments = params.get("arguments", {})
            result = handle_tool_call(tool_name, arguments)
            resp = {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": {
                    "content": [{"type": "text", "text": json.dumps(result, indent=2)}],
                },
            }
        else:
            resp = {
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {"code": -32601, "message": f"Method not found: {method}"},
            }

        print(json.dumps(resp), flush=True)


if __name__ == "__main__":
    mcp_server()