"""MCP (Model Context Protocol) client for SSE-based tool servers.

This module provides a client for connecting to MCP servers via SSE transport,
like those provided by n8n's MCP trigger nodes.

Usage:
    from agent_factory.tools import MCPClient

    # Connect to an SSE MCP server
    client = MCPClient(url="https://your-n8n.cloud/mcp/your-server/sse")
    client.connect()

    # Get tools for the orchestrator
    tools = client.get_tools()

    # Tools are automatically usable with the orchestrator

References:
    - MCP Specification: https://modelcontextprotocol.io/
    - n8n MCP Trigger: https://docs.n8n.io/integrations/builtin/cluster-nodes/sub-nodes/n8n-nodes-langchain.toolmcp/
"""

import json
import threading
import time
from dataclasses import dataclass
from queue import Queue, Empty
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import urlparse, urlunparse

import httpx

from agent_factory.tools.base import BaseTool, ToolResult


@dataclass
class MCPTool:
    """Tool definition from an MCP server."""

    name: str
    description: str
    input_schema: Dict[str, Any]


class MCPToolWrapper(BaseTool):
    """Wrapper that converts an MCP tool to a BaseTool."""

    def __init__(
        self,
        mcp_tool: MCPTool,
        executor: Callable[[str, Dict[str, Any]], Any],
    ):
        """Initialize wrapper.

        Args:
            mcp_tool: The MCP tool definition
            executor: Function to execute the tool (calls the MCP server)
        """
        self._mcp_tool = mcp_tool
        self._executor = executor
        self.name = mcp_tool.name
        self.description = mcp_tool.description

    @property
    def input_schema(self) -> Dict[str, Any]:
        """Return the MCP tool's input schema."""
        return self._mcp_tool.input_schema

    def execute(self, **kwargs) -> Any:
        """Execute the tool by calling the MCP server."""
        return self._executor(self.name, kwargs)


class MCPClient:
    """Client for connecting to MCP servers via SSE transport.

    This client connects to MCP servers that use Server-Sent Events (SSE)
    for communication, such as n8n's MCP trigger nodes.

    The SSE transport works as follows:
    1. Client connects to the SSE endpoint
    2. Server sends an 'endpoint' event with the POST URL
    3. Client sends JSON-RPC requests via POST to that endpoint
    4. Server sends responses via the SSE stream

    Example:
        # Connect to n8n MCP server
        client = MCPClient(
            url="https://your-n8n.cloud/mcp/quickbooks/sse"
        )
        client.connect()

        # Get tools for use with AgentOrchestrator
        tools = client.get_tools()
        for tool in tools:
            orchestrator.register_tool(tool)

        # Execute a tool directly
        result = client.execute_tool("Get_many_invoices", {"Limit": 10})
    """

    PROTOCOL_VERSION = "2024-11-05"

    def __init__(
        self,
        url: str,
        timeout: float = 60.0,
        headers: Optional[Dict[str, str]] = None,
    ):
        """Initialize MCP SSE client.

        Args:
            url: The SSE endpoint URL of the MCP server
            timeout: Request timeout in seconds
            headers: Optional additional headers for requests
        """
        self.url = url
        self.timeout = timeout
        self.headers = headers or {}

        self._post_endpoint: Optional[str] = None
        self._tools: List[MCPTool] = []
        self._message_id = 0
        self._response_map: Dict[int, Queue] = {}
        self._sse_thread: Optional[threading.Thread] = None
        self._connected = False
        self._running = False
        self._http_client: Optional[httpx.Client] = None

    def connect(self) -> None:
        """Connect to the MCP server and fetch tools.

        Raises:
            RuntimeError: If connection fails
        """
        if self._connected:
            return

        self._http_client = httpx.Client(timeout=self.timeout)
        self._running = True

        # Start SSE listener thread
        self._sse_thread = threading.Thread(target=self._listen_sse, daemon=True)
        self._sse_thread.start()

        # Wait for endpoint to be received
        for _ in range(50):  # 5 second timeout
            if self._post_endpoint:
                break
            time.sleep(0.1)

        if not self._post_endpoint:
            self.disconnect()
            raise RuntimeError("Failed to receive endpoint from MCP server")

        # Initialize connection
        self._initialize()

        # Fetch available tools
        self._fetch_tools()

        self._connected = True

    def _listen_sse(self) -> None:
        """Background thread that listens to SSE events."""
        try:
            with httpx.stream(
                "GET",
                self.url,
                headers={**self.headers, "Accept": "text/event-stream"},
                timeout=None,
            ) as response:
                event_type = None
                data_lines: List[str] = []

                for line in response.iter_lines():
                    if not self._running:
                        break

                    line = line.strip()

                    if not line:
                        # Empty line means end of event
                        if event_type and data_lines:
                            data = "\n".join(data_lines)
                            self._handle_sse_event(event_type, data)
                        event_type = None
                        data_lines = []
                        continue

                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                    elif line.startswith("data:"):
                        data_lines.append(line[5:].strip())

        except Exception as e:
            if self._running:
                print(f"SSE connection error: {e}")

    def _handle_sse_event(self, event_type: str, data: str) -> None:
        """Handle an SSE event from the server."""
        if event_type == "endpoint":
            # Server is telling us where to POST messages
            # Handle relative URLs by prepending the base URL
            if data.startswith("/"):
                parsed = urlparse(self.url)
                base_url = f"{parsed.scheme}://{parsed.netloc}"
                self._post_endpoint = base_url + data
            else:
                self._post_endpoint = data
        elif event_type == "message":
            try:
                message = json.loads(data)
                message_id = message.get("id")
                if message_id is not None and message_id in self._response_map:
                    self._response_map[message_id].put(message)
            except json.JSONDecodeError:
                pass

    def _send_request(
        self,
        method: str,
        params: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Send a JSON-RPC request and wait for response.

        Args:
            method: JSON-RPC method name
            params: Optional method parameters

        Returns:
            Response result

        Raises:
            RuntimeError: If request fails or times out
        """
        if not self._post_endpoint:
            raise RuntimeError("Not connected to MCP server")

        self._message_id += 1
        request_id = self._message_id

        # Create response queue for this request
        response_queue: Queue = Queue()
        self._response_map[request_id] = response_queue

        try:
            # Build request
            request = {
                "jsonrpc": "2.0",
                "method": method,
                "id": request_id,
            }
            if params:
                request["params"] = params

            # Send POST request
            response = self._http_client.post(
                self._post_endpoint,
                json=request,
                headers=self.headers,
            )
            response.raise_for_status()

            # Wait for response via SSE
            try:
                result = response_queue.get(timeout=self.timeout)

                if "error" in result:
                    error = result["error"]
                    raise RuntimeError(
                        f"MCP error: {error.get('message', 'Unknown error')}"
                    )

                return result.get("result", {})

            except Empty:
                raise RuntimeError(f"Timeout waiting for response to {method}")

        finally:
            # Clean up response queue
            del self._response_map[request_id]

    def _initialize(self) -> None:
        """Initialize the MCP connection."""
        self._send_request(
            "initialize",
            {
                "protocolVersion": self.PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": {
                    "name": "agent-factory",
                    "version": "0.4.0",
                },
            },
        )

        # Send initialized notification (no response expected)
        if self._post_endpoint and self._http_client:
            self._http_client.post(
                self._post_endpoint,
                json={
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized",
                },
                headers=self.headers,
            )

    def _fetch_tools(self) -> None:
        """Fetch available tools from the server."""
        result = self._send_request("tools/list")
        mcp_tools = result.get("tools", [])

        self._tools = []
        for t in mcp_tools:
            self._tools.append(
                MCPTool(
                    name=t["name"],
                    description=t.get("description", ""),
                    input_schema=t.get("inputSchema", {"type": "object", "properties": {}}),
                )
            )

    def reconnect(self) -> None:
        """Reconnect to the MCP server.

        Useful when the session has expired (e.g., 406 Not Acceptable).
        """
        self.disconnect()
        self.connect()

    def execute_tool(
        self,
        name: str,
        arguments: Optional[Dict[str, Any]] = None,
        _retry: bool = True,
    ) -> Any:
        """Execute a tool on the MCP server.

        Args:
            name: Tool name
            arguments: Tool arguments
            _retry: Internal flag to prevent infinite retry loops

        Returns:
            Tool execution result
        """
        if not self._connected:
            raise RuntimeError("Not connected to MCP server")

        try:
            result = self._send_request(
                "tools/call",
                {
                    "name": name,
                    "arguments": arguments or {},
                },
            )
        except Exception as e:
            # Check if this is a session expiry error (406 Not Acceptable)
            error_str = str(e)
            if _retry and ("406" in error_str or "Not Acceptable" in error_str):
                # Session expired, reconnect and retry once
                self.reconnect()
                return self.execute_tool(name, arguments, _retry=False)
            raise

        # Extract content from MCP response
        content = result.get("content", [])

        if content and isinstance(content, list):
            # MCP returns array of content items
            texts = []
            for item in content:
                if item.get("type") == "text":
                    texts.append(item.get("text", ""))
                elif item.get("type") == "resource":
                    texts.append(json.dumps(item.get("resource", {})))

            if len(texts) == 1:
                # Try to parse as JSON
                try:
                    return json.loads(texts[0])
                except json.JSONDecodeError:
                    return texts[0]
            return "\n".join(texts)

        return result

    def get_tools(self) -> List[MCPToolWrapper]:
        """Get all available tools as BaseTool instances.

        Returns:
            List of tools ready to register with AgentOrchestrator
        """
        return [
            MCPToolWrapper(tool, self.execute_tool)
            for tool in self._tools
        ]

    def get_tool_names(self) -> List[str]:
        """Get names of all available tools.

        Returns:
            List of tool names
        """
        return [t.name for t in self._tools]

    def disconnect(self) -> None:
        """Disconnect from the MCP server."""
        self._running = False
        self._connected = False

        if self._http_client:
            self._http_client.close()
            self._http_client = None

        self._post_endpoint = None
        self._tools = []
        self._response_map.clear()

    def __enter__(self):
        """Context manager entry."""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect()


# Singleton instances for Lambda reuse
_mcp_clients: Dict[str, MCPClient] = {}


def get_mcp_client(
    url: str,
    timeout: float = 60.0,
    headers: Optional[Dict[str, str]] = None,
) -> MCPClient:
    """Get or create an MCP client for a URL.

    This function provides singleton behavior for Lambda environments,
    where connections should be reused across invocations.

    Args:
        url: The SSE endpoint URL of the MCP server
        timeout: Request timeout in seconds
        headers: Optional additional headers

    Returns:
        Connected MCPClient instance
    """
    if url not in _mcp_clients:
        client = MCPClient(url=url, timeout=timeout, headers=headers)
        client.connect()
        _mcp_clients[url] = client

    return _mcp_clients[url]
