"""MCP server for the Cloudflare MCP Gateway using bare mcp library."""

import asyncio
import logging
from typing import Any

import mcp.types as types
from mcp.server.lowlevel import Server, NotificationOptions
from mcp.server.models import InitializationOptions

from .config import GatewaySettings
from .services import ServiceRegistry, ServiceConfig

logger = logging.getLogger(__name__)

# Default timeout for upstream calls (seconds)
UPSTREAM_TIMEOUT = 60.0
SERVICE_DISCOVERY_TIMEOUT = 30.0

# Global state for the gateway
_settings: GatewaySettings | None = None
_registry: ServiceRegistry | None = None
_service_endpoints: dict[str, str] = {}
_service_requires_auth: dict[str, bool] = {}
_cached_tools: list[dict] | None = None


def _get_auth_headers(service_id: str) -> dict[str, str]:
    """Get authorization headers for upstream calls if required."""
    if not _settings:
        raise RuntimeError("Gateway not initialized")

    # Only include auth headers if the service requires it
    if not _service_requires_auth.get(service_id, True):
        return {}

    return {
        "Authorization": f"Bearer {_settings.cloudflare_api_token}",
        "CF-Account-ID": _settings.cloudflare_account_id,
    }


async def _fetch_service_tools(service: ServiceConfig) -> tuple[str, list[dict]]:
    """Fetch tools from a single service.

    Args:
        service: Service configuration.

    Returns:
        Tuple of (service_id, list of tool dicts with prefixed names).
    """
    # Lazy import to avoid startup cost
    from fastmcp import Client
    from fastmcp.client.transports import StreamableHttpTransport

    transport = StreamableHttpTransport(
        url=service.endpoint, headers=_get_auth_headers(service.id)
    )
    client = Client(transport)

    async with client:
        result = await asyncio.wait_for(
            client.list_tools(),
            timeout=SERVICE_DISCOVERY_TIMEOUT,
        )
        tools = result.tools if hasattr(result, "tools") else result

        # Add prefix to tool names and convert to dicts
        prefixed_tools = []
        for tool in tools:
            prefixed_tools.append({
                "name": f"{service.id}_{tool.name}",
                "description": tool.description or f"Call {tool.name} on {service.id}",
                "inputSchema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
                "service_id": service.id,
                "upstream_name": tool.name,
            })

        return (service.id, prefixed_tools)


async def _get_tools() -> list[dict]:
    """Fetch tools from all services in parallel (with caching for session)."""
    global _cached_tools

    if _cached_tools is not None:
        return _cached_tools

    if not _registry:
        raise RuntimeError("Gateway not initialized")

    services = list(_registry.get_enabled_services())
    logger.info(f"Listing tools from {len(services)} services in parallel")

    # Fetch from all services concurrently
    tasks = [_fetch_service_tools(s) for s in services]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    # Aggregate results
    all_tools = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            logger.warning(f"Failed to list tools from {services[i].id}: {result}")
        else:
            service_id, tools = result
            logger.info(f"Got {len(tools)} tools from {service_id}")
            all_tools.extend(tools)

    logger.info(f"Total tools: {len(all_tools)}")
    _cached_tools = all_tools
    return all_tools


# Create the MCP server instance (lightweight)
mcp = Server(
    name="Cloudflare MCP Gateway",
    instructions="A unified gateway for accessing multiple Cloudflare services. "
    "Tools are prefixed with the service name (e.g., docs_, bindings_).",
)


@mcp.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """Handle list_tools request - fetch from all services in parallel."""
    tool_dicts = await _get_tools()

    return [
        types.Tool(
            name=t["name"],
            description=t["description"],
            inputSchema=t["inputSchema"],
        )
        for t in tool_dicts
    ]


@mcp.call_tool()
async def handle_call_tool(
    name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """Handle call_tool request - route to correct service based on prefix."""
    # Lazy import to avoid startup cost
    from fastmcp import Client
    from fastmcp.client.transports import StreamableHttpTransport

    arguments = arguments or {}

    # Parse service prefix: "docs_search" -> service_id="docs", upstream_name="search"
    parts = name.split("_", 1)
    if len(parts) != 2:
        raise ValueError(f"Invalid tool name format: {name}. Expected 'service_toolname'")

    service_id, upstream_name = parts

    # Validate service exists
    endpoint = _service_endpoints.get(service_id)
    if not endpoint:
        available = list(_service_endpoints.keys())
        raise ValueError(f"Unknown service: {service_id}. Available: {available}")

    logger.debug(f"Routing {name} to {service_id}.{upstream_name}")

    # Call upstream
    transport = StreamableHttpTransport(url=endpoint, headers=_get_auth_headers(service_id))
    client = Client(transport)

    try:
        async with client:
            result = await asyncio.wait_for(
                client.call_tool(upstream_name, arguments),
                timeout=UPSTREAM_TIMEOUT,
            )
            logger.debug(f"Upstream {service_id}.{upstream_name} returned {len(result.content)} items")
            return list(result.content)
    except asyncio.TimeoutError:
        logger.error(f"Upstream call to {service_id}.{upstream_name} timed out")
        raise RuntimeError(f"Upstream call timed out after {UPSTREAM_TIMEOUT}s")


def get_initialization_options() -> InitializationOptions:
    """Get initialization options for the server."""
    return InitializationOptions(
        server_name="Cloudflare MCP Gateway",
        server_version="0.1.0",
        capabilities=mcp.get_capabilities(
            notification_options=NotificationOptions(),
            experimental_capabilities={},
        ),
    )


async def initialize_gateway(settings: GatewaySettings) -> None:
    """Initialize the gateway with settings (no connections).

    Args:
        settings: Gateway configuration settings.
    """
    global _settings, _registry, _service_endpoints, _service_requires_auth

    _settings = settings
    _registry = ServiceRegistry(settings)

    # Build endpoints and auth requirement maps
    for service in _registry.get_enabled_services():
        _service_endpoints[service.id] = service.endpoint
        _service_requires_auth[service.id] = service.requires_auth

    logger.info(
        f"Initialized registry with {len(_registry)} services: "
        f"{', '.join(s.id for s in _registry)}"
    )


async def shutdown_gateway() -> None:
    """Perform graceful shutdown of the gateway."""
    logger.info("Shutting down gateway...")
    logger.info("Gateway shutdown complete")
