#!/usr/bin/env python3
"""
gemma4_mcp.py — Local + cloud LLM MCP server.

Two-tier routing:
  1. Local Gemma4-12B (llama.cpp at http://127.0.0.1:8081/v1)
  2. FreeLLMAPI proxy (http://127.0.0.1:3001/v1) — 16 free providers, auto-failover

Set FREELLMAPI_KEY env var, or the proxy's auto-generated key is used.
"""

import json
import os
import sys

try:
    import httpx
except ImportError:
    httpx = None

GEMMA_URL = "http://127.0.0.1:8081/v1/chat/completions"
GEMMA_MODELS_URL = "http://127.0.0.1:8081/v1/models"
FREELLMAPI_URL = "http://127.0.0.1:3001/v1/chat/completions"
FREELLMAPI_KEY = os.environ.get("FREELLMAPI_KEY",
    "freellmapi-69dbee85c7c088c9ea8751a338f85044598bed37e3fe33bf")


def get_model_id() -> str:
    try:
        with httpx.Client(timeout=5) as client:
            resp = client.get(GEMMA_MODELS_URL, headers={"Authorization": "Bearer none"})
            return resp.json()["models"][0]["name"]
    except Exception:
        return "gemma4-12b"


def call_model(url: str, api_key: str, question: str, system: str = "",
               max_tokens: int = 2000, temperature: float = 0.3,
               model: str = "auto", timeout: int = 120) -> dict:
    if httpx is None:
        return {"error": "httpx not installed"}
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": question})
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
    }
    try:
        with httpx.Client(timeout=timeout) as client:
            resp = client.post(url, json=payload, headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            })
            resp.raise_for_status()
            data = resp.json()
        msg = data["choices"][0]["message"]
        content = msg.get("content", "")
        reasoning = msg.get("reasoning_content", "")
        usage = data.get("usage", {})
        timings = data.get("timings", {})
        return {
            "content": content,
            "reasoning": reasoning,
            "tokens_in": usage.get("prompt_tokens", 0),
            "tokens_out": usage.get("completion_tokens", 0),
            "prompt_speed": round(timings.get("prompt_per_second", 0), 1),
            "generation_speed": round(timings.get("predicted_per_second", 0), 1),
        }
    except Exception as e:
        return {"error": str(e)}


def call_gemma(question, system="", max_tokens=2000, temperature=0.3, enable_thinking=True):
    model_id = get_model_id()
    result = call_model(GEMMA_URL, "none", question, system, max_tokens, temperature, model_id)
    if "error" not in result:
        return result
    result = call_model(FREELLMAPI_URL, FREELLMAPI_KEY, question, system,
                        max_tokens, temperature, "auto")
    if "error" not in result:
        return result
    return {"error": "both local Gemma4 and FreeLLMAPI failed"}


def mcp_server():
    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": "gemma4", "version": "0.2.0"},
                },
            }
        elif method == "tools/list":
            resp = {
                "jsonrpc": "2.0", "id": req_id,
                "result": {
                    "tools": [{
                        "name": "gemma4",
                        "description": (
                            "Ask a question to the local Gemma4-12B model, "
                            "with automatic fallback to 16 free LLM providers "
                            "(Gemini, Groq, GPT-4o, etc.). Good at math and code."
                        ),
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "question": {"type": "string", "description": "The question to ask"},
                                "system": {"type": "string", "description": "System prompt (optional)"},
                                "max_tokens": {"type": "number", "description": "Max response tokens (default: 2000)"},
                                "enable_thinking": {"type": "boolean", "description": "Enable reasoning mode (default: true)"},
                            },
                            "required": ["question"],
                        },
                    }],
                },
            }
        elif method == "tools/call":
            tool_name = params.get("name", "")
            arguments = params.get("arguments", {})
            if tool_name == "gemma4":
                result = call_gemma(
                    question=arguments.get("question", ""),
                    system=arguments.get("system", ""),
                    max_tokens=arguments.get("max_tokens", 2000),
                    enable_thinking=arguments.get("enable_thinking", True),
                )
            else:
                result = {"error": f"Unknown tool: {tool_name}"}
            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()