"""
BuildShield MCP Server — stdio transport entry point.

Usage in Claude Desktop / Cursor / any MCP client:

  {
    "mcpServers": {
      "buildshield": {
        "command": "buildshield-mcp",
        "env": { "BUILDSHIELD_API_KEY": "bsk_..." }
      }
    }
  }

Or with npx (npm wrapper — add later):
  "command": "npx", "args": ["-y", "buildshield-mcp"]
"""

from __future__ import annotations

import asyncio
import os
from typing import Any

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp import types

from .tools import (
    handle_check_cost,
    handle_recommend_model,
    handle_get_pricing,
    handle_optimize_prompt,
    handle_check_budget,
    handle_log_usage,
)

app = Server("buildshield-mcp")

# ---------------------------------------------------------------------------
# Tool schemas
# ---------------------------------------------------------------------------

TOOLS: list[types.Tool] = [
    types.Tool(
        name="buildshield_check_cost",
        description=(
            "Calculate the exact USD cost for a specific model and token count. "
            "Works offline — no API key needed. "
            "Use this before making expensive LLM calls to verify cost."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "model":         {"type": "string", "description": "Model name, e.g. 'gpt-4.1', 'claude-sonnet-4.6', 'gemini-2.5-flash'"},
                "input_tokens":  {"type": "integer", "description": "Number of input/prompt tokens"},
                "output_tokens": {"type": "integer", "description": "Number of output/completion tokens (default 0)", "default": 0},
            },
            "required": ["model", "input_tokens"],
        },
    ),
    types.Tool(
        name="buildshield_recommend_model",
        description=(
            "Get the best model recommendations for a task, filtered by cost and context window. "
            "Works offline. Returns a ranked table with pricing."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "task": {"type": "string", "description": "What you're using the model for, e.g. 'code review', 'summarization', 'complex reasoning'"},
                "max_cost_per_1m_input": {"type": "number", "description": "Maximum acceptable input cost per 1M tokens (USD)"},
                "min_context_tokens":    {"type": "integer", "description": "Minimum context window required (tokens)"},
                "provider":              {"type": "string", "description": "Filter to a specific provider: openai, anthropic, google, deepseek"},
            },
            "required": ["task"],
        },
    ),
    types.Tool(
        name="buildshield_get_pricing",
        description=(
            "Get current pricing for one model or browse all models. "
            "Works offline. Returns input/output price per 1M tokens, context window, and tier."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "model":    {"type": "string", "description": "Specific model to look up (optional — omit to see all models)"},
                "provider": {"type": "string", "description": "Filter by provider: openai, anthropic, google, deepseek"},
            },
        },
    ),
    types.Tool(
        name="buildshield_optimize_prompt",
        description=(
            "Compress a prompt using rule-based techniques to reduce token count by 20-40%. "
            "No AI required — works completely offline. "
            "Removes filler phrases, normalizes whitespace, deduplicates repeated instructions."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "The prompt text to optimize"},
                "aggressiveness": {
                    "type": "string",
                    "enum": ["light", "medium", "heavy"],
                    "description": "Compression level: light=whitespace only, medium=filler phrases, heavy=deduplication (default: medium)",
                    "default": "medium",
                },
            },
            "required": ["text"],
        },
    ),
    types.Tool(
        name="buildshield_check_budget",
        description=(
            "Check your remaining BuildShield budget for the current period. "
            "Requires a BuildShield API key (bsk_...) from https://thebuildshield.com/settings."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "api_key": {"type": "string", "description": "Your BuildShield API key (bsk_...). Falls back to BUILDSHIELD_API_KEY env var."},
                "period":  {"type": "string", "enum": ["today", "week", "month"], "description": "Budget period to check (default: today)", "default": "today"},
            },
        },
    ),
    types.Tool(
        name="buildshield_log_usage",
        description=(
            "Log LLM usage to BuildShield for tracking and budget management. "
            "Always calculates cost locally (offline). Syncs to dashboard if api_key provided."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "model":            {"type": "string", "description": "Model used, e.g. 'gpt-4.1'"},
                "input_tokens":     {"type": "integer", "description": "Input token count"},
                "output_tokens":    {"type": "integer", "description": "Output token count"},
                "api_key":          {"type": "string", "description": "BuildShield API key to sync to dashboard (optional)"},
                "task_description": {"type": "string", "description": "Brief description of the task for your records (optional)"},
            },
            "required": ["model", "input_tokens", "output_tokens"],
        },
    ),
]


# ---------------------------------------------------------------------------
# MCP handlers
# ---------------------------------------------------------------------------

@app.list_tools()
async def list_tools() -> list[types.Tool]:
    return TOOLS


@app.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
    try:
        if name == "buildshield_check_cost":
            result = handle_check_cost(
                model=arguments["model"],
                input_tokens=int(arguments["input_tokens"]),
                output_tokens=int(arguments.get("output_tokens", 0)),
            )

        elif name == "buildshield_recommend_model":
            result = handle_recommend_model(
                task=arguments["task"],
                max_cost_per_1m_input=arguments.get("max_cost_per_1m_input"),
                min_context_tokens=arguments.get("min_context_tokens"),
                provider=arguments.get("provider"),
            )

        elif name == "buildshield_get_pricing":
            result = handle_get_pricing(
                model=arguments.get("model"),
                provider=arguments.get("provider"),
            )

        elif name == "buildshield_optimize_prompt":
            result = handle_optimize_prompt(
                text=arguments["text"],
                aggressiveness=arguments.get("aggressiveness", "medium"),
            )

        elif name == "buildshield_check_budget":
            api_key = arguments.get("api_key") or os.environ.get("BUILDSHIELD_API_KEY")
            result = await handle_check_budget(
                api_key=api_key,
                period=arguments.get("period", "today"),
            )

        elif name == "buildshield_log_usage":
            api_key = arguments.get("api_key") or os.environ.get("BUILDSHIELD_API_KEY")
            result = await handle_log_usage(
                model=arguments["model"],
                input_tokens=int(arguments["input_tokens"]),
                output_tokens=int(arguments["output_tokens"]),
                api_key=api_key,
                task_description=arguments.get("task_description"),
            )

        else:
            result = f"Unknown tool: {name}"

    except Exception as e:
        result = f"Error running {name}: {e}"

    return [types.TextContent(type="text", text=result)]


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

async def main() -> None:
    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options(),
        )


def main_sync() -> None:
    asyncio.run(main())


if __name__ == "__main__":
    main_sync()
