"""FastMCP server exposing the platform as MCP tools."""

from __future__ import annotations

import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from fastmcp import FastMCP

from pydanticai_multiagent.config import settings
from pydanticai_multiagent.dependencies import SearchDeps
from pydanticai_multiagent.dependencies.mocks import MockCache, MockHttpClient, MockVectorStore

from .auth import MCPAuthError, validate_api_key

if TYPE_CHECKING:
    from pydanticai_multiagent.services import AgentRegistry, ConversationService, UsageTracker
    from pydanticai_multiagent.services.tenant import Tenant


@dataclass
class MCPState:
    """Shared state for MCP server."""

    db: Any = None
    conversation_service: ConversationService | None = None
    agent_registry: AgentRegistry | None = None
    usage_tracker: UsageTracker | None = None
    http_client: Any = None
    cache: Any = None
    mcp_toolsets: list[Any] | None = None


# Global state holder - set by FastAPI lifespan or CLI
_state = MCPState()


def set_mcp_state(
    db: Any,
    conversation_service: Any,
    agent_registry: Any,
    usage_tracker: Any | None = None,
    http_client: Any | None = None,
    cache: Any | None = None,
    mcp_toolsets: list[Any] | None = None,
) -> None:
    """Set the MCP server state (called by FastAPI lifespan or CLI)."""
    global _state
    _state = MCPState(
        db=db,
        conversation_service=conversation_service,
        agent_registry=agent_registry,
        usage_tracker=usage_tracker,
        http_client=http_client,
        cache=cache,
        mcp_toolsets=mcp_toolsets,
    )


def get_mcp_state() -> MCPState:
    """Get the current MCP state."""
    return _state


@asynccontextmanager
async def mcp_lifespan(server: FastMCP):
    """Lifespan for standalone MCP server (stdio transport)."""
    from pydanticai_multiagent.db import create_pool
    from pydanticai_multiagent.dependencies.mocks import MockDatabase
    from pydanticai_multiagent.services import AgentRegistry, UsageTracker
    from pydanticai_multiagent.services.conversation import ConversationService
    from pydanticai_multiagent.services.message_store import MessageStore

    # Initialize database
    db_url = settings.database_url
    use_real_db = db_url and db_url.startswith("postgresql://")

    if use_real_db:
        db = await create_pool(db_url)
    else:
        db = MockDatabase()

    # Initialize services
    message_store = MessageStore(db)
    usage_tracker = UsageTracker(db)
    conversation_service = ConversationService(message_store, usage_tracker=usage_tracker)

    # Initialize agent registry
    from pydanticai_multiagent import (
        analyst_agent,
        code_agent,
        research_agent,
        router_agent,
        support_agent,
        writer_agent,
    )

    registry = AgentRegistry(db)
    registry.register_builtin("router", router_agent, "Routes requests to specialists")
    registry.register_builtin("research", research_agent, "Web research and information gathering")
    registry.register_builtin("analyst", analyst_agent, "Data analysis and insights")
    registry.register_builtin("code", code_agent, "Code generation and review")
    registry.register_builtin("writer", writer_agent, "Content writing and editing")
    registry.register_builtin("support", support_agent, "Customer support assistance")
    await registry.sync_builtin_agents_to_db()

    # Set global state
    set_mcp_state(
        db=db,
        conversation_service=conversation_service,
        agent_registry=registry,
        usage_tracker=usage_tracker,
    )

    yield {}

    # Cleanup
    if use_real_db and hasattr(db, "close"):
        await db.close()


# Create the FastMCP server
mcp_server = FastMCP(
    name="PydanticAI Platform",
    instructions="""Multi-tenant AI agent platform.

Each tool requires an 'api_key' parameter for authentication.
Use your tenant API key (sk-tenant-xxx format) to authenticate.

Available tools:
- chat: Send a message to an AI agent
- list_agents: List available agents for your tenant
- get_usage: Get your usage statistics
- list_conversations: List your recent conversations
- get_conversation: Get full conversation history
- clear_conversation: Clear a conversation
""",
    lifespan=mcp_lifespan,
)


def _generate_conversation_id() -> str:
    """Generate a unique conversation ID."""
    return f"conv_{uuid.uuid4().hex[:12]}"


async def _get_deps_for_tenant(tenant: Tenant) -> SearchDeps:
    """Create dependencies for a tenant."""
    state = get_mcp_state()
    return SearchDeps(
        http_client=state.http_client or MockHttpClient(),
        db=state.db,
        cache=state.cache or MockCache(),
        user_id=f"tenant:{tenant.id}",
        vector_store=MockVectorStore(),
        search_api_key=settings.search_api_key,
    )


@mcp_server.tool
async def chat(
    api_key: str,
    prompt: str,
    agent: str | None = None,
    conversation_id: str | None = None,
) -> dict[str, Any]:
    """Send a message to an AI agent and get a response.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).
        prompt: The message to send to the agent.
        agent: Agent to use (optional, uses tenant default if not specified).
        conversation_id: Continue an existing conversation (optional).

    Returns:
        Response from the agent including the text, conversation_id, agent name, and model.
    """
    state = get_mcp_state()

    try:
        tenant = await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.agent_registry or not state.conversation_service:
        return {"error": "Platform services not initialized"}

    # Get the agent
    agent_name = agent
    result = None

    if agent_name:
        result = await state.agent_registry.get_agent_for_tenant(agent_name, tenant.id)

    if not result:
        # Try default agent or fall back to first available
        result = await state.agent_registry.get_tenant_default_agent(tenant.id)
        if not result:
            agents = await state.agent_registry.list_agents()
            if not agents:
                return {"error": "No agents available"}
            agent_name = agents[0].name
            result = await state.agent_registry.get_agent_for_tenant(agent_name, tenant.id)

    if not result:
        # Try direct agent lookup (development mode)
        agent_obj = await state.agent_registry.get_agent(agent_name or "router")
        if not agent_obj:
            return {"error": f"Agent '{agent_name}' not available for this tenant"}
        model = tenant.default_model
    else:
        agent_obj, model = result
        if not agent_name:
            # Get the actual agent name from registry
            agents = await state.agent_registry.list_agents()
            agent_name = next((a.name for a in agents if a.name), "default")

    # Generate conversation ID
    conv_id = conversation_id or _generate_conversation_id()

    # Get deps for this tenant
    deps = await _get_deps_for_tenant(tenant)

    # Check token limit
    if tenant.monthly_token_limit and state.usage_tracker:
        within_limit, tokens_used = await state.usage_tracker.check_limit(
            user_id=f"tenant:{tenant.id}",
            token_limit=tenant.monthly_token_limit,
            period_days=30,
        )
        if not within_limit:
            return {
                "error": "Monthly token limit exceeded",
                "tokens_used": tokens_used,
                "limit": tenant.monthly_token_limit,
            }

    # Run the agent
    try:
        response = await state.conversation_service.chat_sync(
            conversation_id=conv_id,
            user_id=f"tenant:{tenant.id}",
            prompt=prompt,
            agent=agent_obj,
            deps=deps,
            agent_name=agent_name or "default",
            model=model,
            toolsets=state.mcp_toolsets or [],
        )

        return {
            "response": response,
            "conversation_id": conv_id,
            "agent": agent_name or "default",
            "model": model,
        }
    except Exception as e:
        return {"error": f"Agent error: {e!s}"}


@mcp_server.tool
async def list_agents(api_key: str) -> dict[str, Any]:
    """List all agents available to your tenant.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).

    Returns:
        List of available agents with their names and descriptions.
    """
    state = get_mcp_state()

    try:
        tenant = await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.agent_registry:
        return {"error": "Platform services not initialized"}

    # Try to get tenant-specific agents
    from pydanticai_multiagent.services import TenantService

    if state.db:
        tenant_service = TenantService(state.db)
        agents = await tenant_service.get_tenant_agents(tenant.id)
        if agents:
            return {
                "agents": [
                    {"name": a["name"], "description": a["description"]} for a in agents
                ]
            }

    # Fall back to all registered agents
    all_agents = await state.agent_registry.list_agents()
    return {
        "agents": [
            {"name": a.name, "description": a.description} for a in all_agents
        ]
    }


@mcp_server.tool
async def get_usage(api_key: str, days: int = 30) -> dict[str, Any]:
    """Get usage statistics for your tenant.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).
        days: Number of days to include (default 30).

    Returns:
        Usage statistics including total requests, tokens, and cost breakdown.
    """
    state = get_mcp_state()

    try:
        tenant = await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.usage_tracker:
        return {"error": "Usage tracking not available"}

    summary = await state.usage_tracker.get_user_summary(f"tenant:{tenant.id}", days)

    return {
        "tenant_id": tenant.id,
        "period_days": days,
        "total_requests": summary.total_requests,
        "total_tokens": summary.total_tokens,
        "estimated_cost_usd": summary.estimated_cost,
        "by_model": summary.by_model,
        "by_agent": summary.by_agent,
    }


@mcp_server.tool
async def list_conversations(api_key: str, limit: int = 20) -> dict[str, Any]:
    """List recent conversations for your tenant.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).
        limit: Maximum number of conversations to return (default 20).

    Returns:
        List of conversation summaries with IDs and metadata.
    """
    state = get_mcp_state()

    try:
        tenant = await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.conversation_service:
        return {"error": "Conversation service not available"}

    conversations = await state.conversation_service.list_conversations(
        user_id=f"tenant:{tenant.id}",
        limit=limit,
    )

    return {
        "conversations": [
            {
                "conversation_id": c.get("conversation_id", ""),
                "created_at": str(c.get("created_at")) if c.get("created_at") else None,
                "updated_at": str(c.get("updated_at")) if c.get("updated_at") else None,
                "message_count": c.get("message_count", 0),
            }
            for c in conversations
        ]
    }


@mcp_server.tool
async def get_conversation(api_key: str, conversation_id: str) -> dict[str, Any]:
    """Get the full history of a conversation.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).
        conversation_id: The conversation ID to retrieve.

    Returns:
        Full conversation history with all messages.
    """
    state = get_mcp_state()

    try:
        tenant = await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.conversation_service:
        return {"error": "Conversation service not available"}

    context = await state.conversation_service.get_context(
        conversation_id=conversation_id,
        user_id=f"tenant:{tenant.id}",
    )

    # Convert messages to dicts
    messages = []
    for msg in context.messages:
        if hasattr(msg, "model_dump"):
            messages.append(msg.model_dump())
        else:
            messages.append({"content": str(msg)})

    return {
        "conversation_id": conversation_id,
        "messages": messages,
    }


@mcp_server.tool
async def clear_conversation(api_key: str, conversation_id: str) -> dict[str, Any]:
    """Clear a conversation's history.

    Args:
        api_key: Your tenant API key (sk-tenant-xxx format).
        conversation_id: The conversation ID to clear.

    Returns:
        Status indicating whether the conversation was cleared.
    """
    state = get_mcp_state()

    try:
        await validate_api_key(api_key, state.db)
    except MCPAuthError as e:
        return {"error": e.message}

    if not state.conversation_service:
        return {"error": "Conversation service not available"}

    deleted = await state.conversation_service.clear_conversation(conversation_id)

    return {
        "status": "cleared" if deleted else "not_found",
        "conversation_id": conversation_id,
    }
