"""OpenRouter MCP Server.

Exposes OpenRouter API capabilities as MCP tools for text completion,
image generation, and model discovery.
"""

import base64
import os
from pathlib import Path
from typing import Optional

from dotenv import load_dotenv
from fastmcp import FastMCP
from fastmcp.utilities.types import Image

from mcp_openrouter.client import OpenRouterClient
from mcp_openrouter.config import get_default_model


def _load_env_files() -> None:
    """Load environment variables from the working tree and current directory."""
    package_root = Path(__file__).resolve().parents[2]
    env_paths = (Path.cwd() / ".env", package_root / ".env")
    seen_paths: set[Path] = set()

    for env_path in env_paths:
        resolved_path = env_path.resolve()
        if resolved_path in seen_paths:
            continue
        seen_paths.add(resolved_path)
        load_dotenv(resolved_path)


_load_env_files()

# Initialize FastMCP server
mcp = FastMCP(
    "openrouter",
    instructions="""OpenRouter MCP Server provides access to 300+ AI models.

Available tools:
- chat: Text completion with any model
- generate_image: Image generation with image models
- embed: Generate embeddings with embedding models
- list_models: List available models by capability
- find_models: Search for models by name

Requires OPENROUTER_API_KEY environment variable.""",
)


def _register_tool(fn, *, name: str):
    """Register a FastMCP tool while preserving a directly callable function."""
    mcp.tool(fn, name=name)
    return fn


def get_client() -> OpenRouterClient:
    """Get an initialized OpenRouter client."""
    api_key = os.environ.get("OPENROUTER_API_KEY")
    if not api_key:
        raise ValueError(
            "OPENROUTER_API_KEY environment variable not set. "
            "Get your key at: https://openrouter.ai/keys"
        )
    return OpenRouterClient(api_key)


def _chat(
    prompt: Optional[str] = None,
    messages: Optional[list[dict]] = None,
    model: Optional[str] = None,
    system: Optional[str] = None,
    max_tokens: Optional[int] = None,
    temperature: Optional[float] = None,
    top_p: Optional[float] = None,
    top_k: Optional[int] = None,
    frequency_penalty: Optional[float] = None,
    presence_penalty: Optional[float] = None,
    seed: Optional[int] = None,
    stop: Optional[list[str]] = None,
    json_mode: bool = False,
    response_format: Optional[dict] = None,
    reasoning_effort: Optional[str] = None,
    provider: Optional[dict] = None,
    assistant_prefill: Optional[str] = None,
) -> str:
    """Send a chat completion request to any OpenRouter model.

    Args:
        prompt: User message to send (provide either prompt or messages, not both)
        messages: Multi-turn conversation as a list of {role, content} dicts
            (provide either prompt or messages, not both)
        model: Model identifier (e.g., "anthropic/claude-sonnet-4", "openai/gpt-4o").
            If not specified, uses DEFAULT_TEXT_MODEL environment variable.
        system: Optional system prompt to set context
        max_tokens: Maximum tokens in response (model default if not specified)
        temperature: Sampling temperature 0-2 (model default if not specified)
        top_p: Nucleus sampling threshold 0-1
        top_k: Top-k sampling (number of top tokens to consider)
        frequency_penalty: Penalize repeated tokens (-2 to 2)
        presence_penalty: Penalize tokens already present (-2 to 2)
        seed: Random seed for deterministic outputs
        stop: List of stop sequences
        json_mode: If True, request JSON-formatted response (backward compat)
        response_format: Response format spec, e.g. {"type": "json_schema", ...}.
            Supersedes json_mode if both provided.
        reasoning_effort: Reasoning effort level: "minimal", "medium", or "high"
        provider: Provider routing control (e.g., {"order": ["Anthropic", "Google"]})
        assistant_prefill: Text to prefill the assistant response with

    Returns:
        The model's response text
    """
    if prompt and messages:
        raise ValueError("Provide either 'prompt' or 'messages', not both.")
    if not prompt and not messages:
        raise ValueError("Either 'prompt' or 'messages' must be provided.")

    resolved_model = model or get_default_model("text")
    if not resolved_model:
        raise ValueError(
            "No model specified. Either pass the 'model' parameter or set "
            "DEFAULT_TEXT_MODEL environment variable."
        )

    client = get_client()

    # Build messages list
    msg_list = []
    if messages:
        if system:
            msg_list.append({"role": "system", "content": system})
        msg_list.extend(messages)
    else:
        if system:
            msg_list.append({"role": "system", "content": system})
        msg_list.append({"role": "user", "content": prompt})

    if assistant_prefill:
        msg_list.append({"role": "assistant", "content": assistant_prefill})

    # Build kwargs
    kwargs = {}
    if max_tokens is not None:
        kwargs["max_tokens"] = max_tokens
    if temperature is not None:
        kwargs["temperature"] = temperature
    if top_p is not None:
        kwargs["top_p"] = top_p
    if top_k is not None:
        kwargs["top_k"] = top_k
    if frequency_penalty is not None:
        kwargs["frequency_penalty"] = frequency_penalty
    if presence_penalty is not None:
        kwargs["presence_penalty"] = presence_penalty
    if seed is not None:
        kwargs["seed"] = seed
    if stop is not None:
        kwargs["stop"] = stop
    if reasoning_effort is not None:
        kwargs["reasoning"] = {"effort": reasoning_effort}
    if provider is not None:
        kwargs["provider"] = provider

    # Resolve response format: explicit response_format > json_mode > none
    if response_format is not None:
        kwargs["response_format"] = response_format
    elif json_mode:
        kwargs["response_format"] = {"type": "json_object"}

    result = client.chat(resolved_model, msg_list, **kwargs)
    return result["choices"][0]["message"]["content"]


chat = _register_tool(_chat, name="chat")


def _generate_image(
    prompt: str,
    model: Optional[str] = None,
    aspect_ratio: str = "1:1",
    size: str = "1K",
    background: Optional[str] = None,
    quality: Optional[str] = None,
    output_format: Optional[str] = None,
    output_path: Optional[str] = None,
) -> Image:
    """Generate an image using an OpenRouter image generation model.

    Args:
        prompt: Image description - be specific about style, colors, composition
        model: Image model (e.g., "google/gemini-3-pro-image-preview").
            If not specified, uses DEFAULT_IMAGE_MODEL environment variable.
        aspect_ratio: Aspect ratio - 1:1, 16:9, 9:16, 4:3, 3:4, or 21:9
        size: Image size - 1K, 2K, or 4K
        background: Background setting (e.g., "transparent" for PNG/WebP)
        quality: Quality setting (e.g., "high", "medium", "low")
        output_format: Output format (e.g., "png", "webp", "jpeg")
        output_path: Optional absolute file path to save the image.
            Must be an absolute path if provided (e.g., "/Users/name/output.png").

    Returns:
        The generated image data
    """
    resolved_model = model or get_default_model("image")
    if not resolved_model:
        raise ValueError(
            "No model specified. Either pass the 'model' parameter or set "
            "DEFAULT_IMAGE_MODEL environment variable."
        )

    client = get_client()

    images = client.generate_image(
        resolved_model,
        prompt,
        aspect_ratio=aspect_ratio,
        size=size,
        background=background,
        quality=quality,
        output_format=output_format,
    )

    if not images:
        raise ValueError("No image was generated. Try adjusting the prompt or model.")

    # Extract base64 data from the response
    data_url = images[0]["image_url"]["url"]
    base64_data = data_url.split(",")[1]

    # Determine format from data URL (e.g., "data:image/png;base64,...")
    mime_type = data_url.split(";")[0].split(":")[1]
    img_format = mime_type.split("/")[1]  # "png", "webp", etc.

    # Decode image data
    image_data = base64.b64decode(base64_data)

    # Save to file if output_path is provided
    if output_path:
        path = Path(output_path)
        if not path.is_absolute():
            raise ValueError(
                f"output_path must be an absolute path, got: {output_path}"
            )
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(image_data)

    return Image(data=image_data, format=img_format)


generate_image = _register_tool(_generate_image, name="generate_image")


def _embed(
    input: str | list[str],
    model: Optional[str] = None,
    encoding_format: Optional[str] = None,
    dimensions: Optional[int] = None,
) -> dict:
    """Generate embeddings for text input using an OpenRouter embedding model.

    Args:
        input: Text string or list of strings to embed
        model: Embedding model (e.g., "mistralai/mistral-embed-2312").
            If not specified, uses DEFAULT_EMBEDDING_MODEL environment variable.
        encoding_format: Output format: "float" or "base64" (default: float)
        dimensions: Custom embedding dimensions (model-dependent)

    Returns:
        Embedding response with data array containing embedding vectors and usage info
    """
    resolved_model = model or get_default_model("embedding")
    if not resolved_model:
        raise ValueError(
            "No model specified. Either pass the 'model' parameter or set "
            "DEFAULT_EMBEDDING_MODEL environment variable."
        )

    client = get_client()

    kwargs = {}
    if encoding_format is not None:
        kwargs["encoding_format"] = encoding_format
    if dimensions is not None:
        kwargs["dimensions"] = dimensions

    return client.embeddings(resolved_model, input, **kwargs)


embed = _register_tool(_embed, name="embed")


def _list_models(capability: Optional[str] = None) -> list[dict]:
    """List available OpenRouter models, optionally filtered by capability.

    Args:
        capability: Filter by capability:
            - "vision": Models that can analyze images
            - "image_gen": Models that can generate images
            - "embedding": Models that can generate embeddings
            - "tools": Models that support tool/function calling
            - "long_context": Models with 100k+ context window

    Returns:
        List of models with slug, name, context_length, and pricing info
    """
    client = get_client()
    models = client.list_models(capability)

    # Return simplified model info
    return [
        {
            "slug": m["slug"],
            "name": m["name"],
            "context_length": m.get("context_length"),
            "pricing": {
                "prompt": m.get("pricing", {}).get("prompt"),
                "completion": m.get("pricing", {}).get("completion"),
            },
        }
        for m in models
    ]


list_models = _register_tool(_list_models, name="list_models")


def _find_models(search_term: str) -> list[dict]:
    """Search for models by name or slug.

    Args:
        search_term: Text to search for in model names (e.g., "claude", "gpt", "gemini")

    Returns:
        List of matching models (max 20) with slug, name, and context_length
    """
    client = get_client()
    matches = client.find_model(search_term)

    # Return simplified model info, limited to 20 results
    return [
        {
            "slug": m["slug"],
            "name": m["name"],
            "context_length": m.get("context_length"),
        }
        for m in matches[:20]
    ]


find_models = _register_tool(_find_models, name="find_models")


def main():
    """Run the MCP server."""
    mcp.run()


if __name__ == "__main__":
    main()
