"""
Agent Status MCP Server implementation.

Implements the Model Context Protocol (MCP) for AI assistant integration.
"""

from __future__ import annotations

import json
import sys
import os
import logging
from typing import Any, Dict, List, Optional

try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    Server = object

import agent_status
from agent_status.client import AgentStatusClient, AgentStatusError

logger = logging.getLogger(__name__)


def create_mcp_server(
    api_key: Optional[str] = None,
    base_url: str = "https://api.fabric.carmel.so",
) -> "Server":
    """
    Create an MCP server with Agent Status tools.

    Args:
        api_key: API key. If not provided, uses RORA_API_KEY env var.
        base_url: API base URL

    Returns:
        Configured MCP Server instance

    Example:
        from agent_status.mcp import create_mcp_server
        server = create_mcp_server(api_key="fb_live_xxx")
    """
    if not MCP_AVAILABLE:
        raise ImportError(
            "MCP is not installed. Install with: pip install mcp"
        )

    api_key = api_key or os.environ.get("RORA_API_KEY")
    if not api_key:
        raise ValueError("API key required. Set RORA_API_KEY or pass api_key parameter.")

    client = AgentStatusClient(api_key=api_key, base_url=base_url)
    server = Server("agent-status-mcp")

    @server.list_tools()
    async def list_tools() -> List[Tool]:
        """List available Agent Status tools."""
        return [
            Tool(
                name="agent_status_register",
                description="Register an AI agent for outside-in monitoring. Returns the agent ID.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "name": {
                            "type": "string",
                            "description": "Name for this agent"
                        },
                        "endpoint": {
                            "type": "string",
                            "description": "HTTP endpoint URL for the agent"
                        },
                        "description": {
                            "type": "string",
                            "description": "Optional description of what this agent does"
                        },
                        "tags": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Optional tags for categorization"
                        }
                    },
                    "required": ["name", "endpoint"]
                }
            ),
            Tool(
                name="agent_status_status",
                description="Get the current status and metrics for a registered agent.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "agent_id": {
                            "type": "string",
                            "description": "The agent ID to check"
                        }
                    },
                    "required": ["agent_id"]
                }
            ),
            Tool(
                name="agent_status_run",
                description="Run a one-off validation probe against an agent.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "agent_id": {
                            "type": "string",
                            "description": "The agent ID to probe"
                        },
                        "prompt": {
                            "type": "string",
                            "description": "Optional custom prompt for the probe"
                        }
                    },
                    "required": ["agent_id"]
                }
            ),
            Tool(
                name="agent_status_list",
                description="List all registered agents and their current status.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "limit": {
                            "type": "integer",
                            "description": "Maximum number of agents to return (default: 50)"
                        }
                    }
                }
            ),
            Tool(
                name="agent_status_delete",
                description="Delete a registered agent from monitoring.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "agent_id": {
                            "type": "string",
                            "description": "The agent ID to delete"
                        }
                    },
                    "required": ["agent_id"]
                }
            ),
        ]

    @server.call_tool()
    async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
        """Handle tool calls."""
        try:
            if name == "agent_status_register":
                result = await _handle_register(client, arguments)
            elif name == "agent_status_status":
                result = await _handle_status(client, arguments)
            elif name == "agent_status_run":
                result = await _handle_run(client, arguments)
            elif name == "agent_status_list":
                result = await _handle_list(client, arguments)
            elif name == "agent_status_delete":
                result = await _handle_delete(client, arguments)
            else:
                result = {"error": f"Unknown tool: {name}"}

            return [TextContent(
                type="text",
                text=json.dumps(result, indent=2)
            )]

        except AgentStatusError as e:
            return [TextContent(
                type="text",
                text=json.dumps({"error": str(e)}, indent=2)
            )]
        except Exception as e:
            logger.error(f"Tool error: {e}")
            return [TextContent(
                type="text",
                text=json.dumps({"error": f"Internal error: {str(e)}"}, indent=2)
            )]

    return server


async def _handle_register(client: AgentStatusClient, args: Dict[str, Any]) -> Dict[str, Any]:
    """Handle agent_status_register tool call."""
    result = client.register(
        name=args["name"],
        endpoint=args["endpoint"],
        description=args.get("description"),
        tags=args.get("tags"),
    )
    return {
        "success": True,
        "agent_id": result.agent_id,
        "message": f"Agent '{args['name']}' registered successfully"
    }


async def _handle_status(client: AgentStatusClient, args: Dict[str, Any]) -> Dict[str, Any]:
    """Handle agent_status_status tool call."""
    status = client.status(args["agent_id"])
    return {
        "agent_id": args["agent_id"],
        "verdict": status.verdict.value,
        "uptime_24h": status.uptime_24h,
        "latency_p95_ms": status.latency_p95,
        "pass_rate": status.pass_rate,
        "last_checked": status.last_checked.isoformat() if status.last_checked else None,
    }


async def _handle_run(client: AgentStatusClient, args: Dict[str, Any]) -> Dict[str, Any]:
    """Handle agent_status_run tool call."""
    result = client.run(
        agent_id=args["agent_id"],
        prompt=args.get("prompt"),
    )
    return {
        "agent_id": args["agent_id"],
        "run_id": result.run_id,
        "verdict": result.verdict.value,
        "latency_ms": result.latency_ms,
        "response_preview": result.response_preview,
    }


async def _handle_list(client: AgentStatusClient, args: Dict[str, Any]) -> Dict[str, Any]:
    """Handle agent_status_list tool call."""
    agents = client.list_agents(limit=args.get("limit", 50))
    return {
        "count": len(agents),
        "agents": [
            {
                "id": a.id,
                "name": a.name,
                "status": a.status,
                "last_verdict": a.last_verdict.value if a.last_verdict else None,
            }
            for a in agents
        ]
    }


async def _handle_delete(client: AgentStatusClient, args: Dict[str, Any]) -> Dict[str, Any]:
    """Handle agent_status_delete tool call."""
    client.delete(args["agent_id"])
    return {
        "success": True,
        "agent_id": args["agent_id"],
        "message": "Agent deleted successfully"
    }


def run_mcp_server(
    api_key: Optional[str] = None,
    base_url: str = "https://api.fabric.carmel.so",
) -> None:
    """
    Run the Agent Status MCP server.

    Args:
        api_key: API key
        base_url: API base URL
    """
    if not MCP_AVAILABLE:
        print("Error: MCP is not installed. Install with: pip install mcp", file=sys.stderr)
        sys.exit(1)

    import asyncio

    async def main():
        server = create_mcp_server(api_key=api_key, base_url=base_url)
        async with stdio_server() as (read_stream, write_stream):
            await server.run(read_stream, write_stream)

    asyncio.run(main())
