from fastmcp import FastMCP
import requests
import os
import json
import logging
from typing import List, Dict, Any

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize the server
mcp = FastMCP("OpenRouter Swarm")

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"

MODELS = {
    "deepseek-r1": "deepseek/deepseek-r1",
    "kimi": "moonshot/kimi-k1.5-preview",
    "minimax": "minimax/minimax-01"
}

def call_model(model_key: str, prompt: str, system_prompt: str = "") -> str:
    """Calls a specific model via OpenRouter."""
    if not OPENROUTER_API_KEY:
        return "Error: OPENROUTER_API_KEY not found."
    
    model_id = MODELS.get(model_key, model_key)
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    
    data = {
        "model": model_id,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]
    }
    
    try:
        response = requests.post(OPENROUTER_URL, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        return result['choices'][0]['message']['content']
    except Exception as e:
        logger.error(f"Error calling {model_key}: {e}")
        return f"Error: {str(e)}"

@mcp.tool()
def deepseek_reason(objective: str) -> str:
    """
    Uses DeepSeek R1 (Reasoning Model) to architect a solution.
    """
    system_prompt = "You are a Senior Solutions Architect. Analyze the objective, identify constraints, and output a detailed technical architecture and implementation plan. Use <thinking> tags for your reasoning process."
    return call_model("deepseek-r1", objective, system_prompt)

@mcp.tool()
def kimi_coder(task: str) -> str:
    """
    Uses Kimi k1.5 (Coding Specialist) to generate production-ready code.
    """
    system_prompt = "You are an Expert Python/JS Developer. Write clean, efficient, and documented code for the given task. Output ONLY the code block."
    return call_model("kimi", task, system_prompt)

@mcp.tool()
def minimax_creative(task: str) -> str:
    """
    Uses MiniMax 01 (Creative/Content) for UI copy, marketing text, or creative ideation.
    """
    system_prompt = "You are a Creative Director. Generate engaging copy, UI text, or marketing content."
    return call_model("minimax", task, system_prompt)

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