# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
MCP Client Implementation

Low-level client for communicating with MCP servers over various transports.
Implements the Model Context Protocol specification.

This module provides async communication with MCP servers, handling:
- Connection lifecycle (connect, disconnect, reconnect)
- Message serialization (JSON-RPC 2.0)
- Request/response correlation
- Error handling and timeouts
- Streaming support
"""

import asyncio
import json
import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from .config import MCPServerConfig, MCPTransport

logger = logging.getLogger(__name__)


class MCPClientError(Exception):
    """Base exception for MCP client errors."""

    pass


class MCPConnectionError(MCPClientError):
    """Connection-related errors."""

    pass


class MCPTimeoutError(MCPClientError):
    """Timeout errors."""

    pass


class MCPProtocolError(MCPClientError):
    """Protocol/message format errors."""

    pass


class MCPToolError(MCPClientError):
    """Tool execution errors."""

    def __init__(self, message: str, tool_name: str = None, details: Any = None):
        super().__init__(message)
        self.tool_name = tool_name
        self.details = details


# ============================================================
# MCP DATA STRUCTURES
# ============================================================


@dataclass
class MCPTool:
    """Represents an MCP tool definition."""

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

    # Optional metadata
    annotations: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "inputSchema": self.input_schema,
        }


@dataclass
class MCPResource:
    """Represents an MCP resource."""

    uri: str
    name: str
    description: Optional[str] = None
    mime_type: Optional[str] = None


@dataclass
class MCPPrompt:
    """Represents an MCP prompt template."""

    name: str
    description: Optional[str] = None
    arguments: List[Dict[str, Any]] = field(default_factory=list)


@dataclass
class MCPToolResult:
    """Result from an MCP tool call."""

    content: List[Dict[str, Any]]  # Array of content blocks
    is_error: bool = False

    @property
    def text(self) -> str:
        """Get combined text content."""
        texts = []
        for block in self.content:
            if block.get("type") == "text":
                texts.append(block.get("text", ""))
        return "\n".join(texts)

    @property
    def first_text(self) -> Optional[str]:
        """Get first text content block."""
        for block in self.content:
            if block.get("type") == "text":
                return block.get("text")
        return None


# ============================================================
# TRANSPORT LAYER
# ============================================================


class MCPTransportBase(ABC):
    """Abstract base class for MCP transports."""

    @abstractmethod
    async def connect(self) -> None:
        """Establish connection."""
        pass

    @abstractmethod
    async def disconnect(self) -> None:
        """Close connection."""
        pass

    @abstractmethod
    async def send(self, message: Dict[str, Any]) -> None:
        """Send a message."""
        pass

    @abstractmethod
    async def receive(self) -> Dict[str, Any]:
        """Receive a message."""
        pass

    @property
    @abstractmethod
    def is_connected(self) -> bool:
        """Check if connected."""
        pass


class StdioTransport(MCPTransportBase):
    """
    MCP transport over stdin/stdout of a subprocess.

    This is the most common transport for local MCP servers.
    """

    def __init__(
        self,
        command: str,
        args: List[str],
        env: Optional[Dict[str, str]] = None,
        cwd: Optional[str] = None,
        timeout: float = 30.0,
    ):
        self.command = command
        self.args = args
        self.env = env or os.environ.copy()
        self.cwd = cwd
        self.timeout = timeout

        self._process: Optional[asyncio.subprocess.Process] = None
        self._reader_task: Optional[asyncio.Task] = None
        self._message_queue: asyncio.Queue = asyncio.Queue()
        self._connected = False

    @property
    def is_connected(self) -> bool:
        return self._connected and self._process is not None

    async def connect(self) -> None:
        """Start the subprocess and establish communication."""
        if self._connected:
            return

        try:
            logger.info(f"Starting MCP server: {self.command} {' '.join(self.args)}")

            self._process = await asyncio.create_subprocess_exec(
                self.command,
                *self.args,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=self.env,
                cwd=self.cwd,
            )

            # Start background reader
            self._reader_task = asyncio.create_task(self._read_loop())
            self._connected = True

            logger.info(f"MCP server started (PID: {self._process.pid})")

        except FileNotFoundError:
            raise MCPConnectionError(
                f"Command not found: {self.command}. Make sure the MCP server is installed."
            )
        except Exception as e:
            raise MCPConnectionError(f"Failed to start MCP server: {e}")

    async def disconnect(self) -> None:
        """Stop the subprocess."""
        if not self._connected:
            return

        self._connected = False

        if self._reader_task:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass

        if self._process:
            try:
                self._process.terminate()
                await asyncio.wait_for(self._process.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                self._process.kill()
                await self._process.wait()
            except ProcessLookupError:
                pass  # Already terminated

            self._process = None

        logger.info("MCP server stopped")

    async def send(self, message: Dict[str, Any]) -> None:
        """Send a JSON-RPC message to the server."""
        if not self.is_connected or not self._process.stdin:
            raise MCPConnectionError("Not connected")

        try:
            data = json.dumps(message) + "\n"
            self._process.stdin.write(data.encode())
            await self._process.stdin.drain()
            logger.debug(f"Sent: {message.get('method', message.get('id', 'unknown'))}")
        except Exception as e:
            raise MCPConnectionError(f"Failed to send message: {e}")

    async def receive(self) -> Dict[str, Any]:
        """Receive a JSON-RPC message from the queue."""
        try:
            message = await asyncio.wait_for(self._message_queue.get(), timeout=self.timeout)
            return message
        except asyncio.TimeoutError:
            raise MCPTimeoutError("Timeout waiting for response")

    async def _read_loop(self) -> None:
        """Background task to read messages from stdout."""
        buffer = ""

        while self._connected and self._process and self._process.stdout:
            try:
                # Read a chunk of data
                chunk = await self._process.stdout.read(8192)
                if not chunk:
                    # EOF - process terminated
                    logger.warning("MCP server stdout closed")
                    break

                buffer += chunk.decode()

                # Process complete lines
                while "\n" in buffer:
                    line, buffer = buffer.split("\n", 1)
                    line = line.strip()
                    if not line:
                        continue

                    try:
                        message = json.loads(line)
                        await self._message_queue.put(message)
                        logger.debug(
                            f"Received: {message.get('method', message.get('id', 'notification'))}"
                        )
                    except json.JSONDecodeError as e:
                        logger.warning(f"Invalid JSON from MCP server: {e}")

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error reading from MCP server: {e}")
                break

        # Read stderr for error messages
        if self._process and self._process.stderr:
            try:
                stderr = await self._process.stderr.read()
                if stderr:
                    logger.warning(f"MCP server stderr: {stderr.decode()}")
            except Exception:
                pass


class HTTPTransport(MCPTransportBase):
    """
    MCP transport over HTTP (Streamable HTTP transport).

    Uses HTTP POST for requests and SSE for streaming responses.
    """

    def __init__(
        self,
        url: str,
        headers: Optional[Dict[str, str]] = None,
        timeout: float = 30.0,
    ):
        self.url = url.rstrip("/")
        self.headers = headers or {}
        self.timeout = timeout

        self._session = None
        self._connected = False

    @property
    def is_connected(self) -> bool:
        return self._connected

    async def connect(self) -> None:
        """Initialize HTTP session."""
        try:
            import aiohttp

            self._session = aiohttp.ClientSession(
                headers={
                    "Content-Type": "application/json",
                    **self.headers,
                },
                timeout=aiohttp.ClientTimeout(total=self.timeout),
            )
            self._connected = True
            logger.info(f"Connected to MCP server at {self.url}")
        except ImportError:
            raise MCPConnectionError("aiohttp is required for HTTP transport.")

    async def disconnect(self) -> None:
        """Close HTTP session."""
        if self._session:
            await self._session.close()
            self._session = None
        self._connected = False

    async def send(self, message: Dict[str, Any]) -> None:
        """Send is combined with receive for HTTP."""
        # For HTTP, we don't send separately - it's request/response
        pass

    async def receive(self) -> Dict[str, Any]:
        """For HTTP, receive is handled in request method."""
        raise NotImplementedError("Use request() for HTTP transport")

    async def request(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Send request and receive response."""
        if not self._session:
            raise MCPConnectionError("Not connected")

        try:
            async with self._session.post(
                f"{self.url}/message",
                json=message,
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise MCPProtocolError(f"HTTP {response.status}: {text}")
                return await response.json()
        except Exception as e:
            raise MCPConnectionError(f"HTTP request failed: {e}")


# ============================================================
# MCP CLIENT
# ============================================================


class MCPClient:
    """
    High-level MCP client for communicating with MCP servers.

    Handles the MCP protocol including:
    - Protocol initialization
    - Tool listing and calling
    - Resource listing and reading
    - Prompt listing and getting

    Usage:
        config = MCPServerConfig(
            name="github",
            transport="stdio",
            command="npx",
            args=["-y", "@modelcontextprotocol/server-github"],
        )

        client = MCPClient(config)
        await client.connect()

        tools = await client.list_tools()
        result = await client.call_tool("create_issue", {"title": "Bug"})

        await client.disconnect()
    """

    PROTOCOL_VERSION = "2025-11-05"  # Must match server version

    def __init__(self, config: MCPServerConfig):
        self.config = config
        self.name = config.name

        # Create transport
        if config.transport == MCPTransport.STDIO:
            self._transport = StdioTransport(
                command=config.command,
                args=config.args,
                env=config.get_env(),
                cwd=config.cwd,
                timeout=config.timeout,
            )
        elif config.transport == MCPTransport.HTTP:
            self._transport = HTTPTransport(
                url=config.url,
                headers=config.headers,
                timeout=config.timeout,
            )
        else:
            raise ValueError(f"Unsupported transport: {config.transport}")

        # State
        self._initialized = False
        self._request_id = 0
        self._pending_requests: Dict[int, asyncio.Future] = {}
        self._server_info: Optional[Dict[str, Any]] = None
        self._capabilities: Dict[str, Any] = {}

        # Cached data
        self._tools: Optional[List[MCPTool]] = None
        self._resources: Optional[List[MCPResource]] = None
        self._prompts: Optional[List[MCPPrompt]] = None

    @property
    def is_connected(self) -> bool:
        return self._transport.is_connected and self._initialized

    @property
    def server_info(self) -> Optional[Dict[str, Any]]:
        return self._server_info

    @property
    def capabilities(self) -> Dict[str, Any]:
        return self._capabilities

    async def connect(self) -> None:
        """Connect to the MCP server and initialize protocol."""
        if self._initialized:
            return

        # Connect transport
        await self._transport.connect()

        # Start message handler for stdio transport
        if isinstance(self._transport, StdioTransport):
            asyncio.create_task(self._message_handler())

        # Send initialize request
        try:
            response = await self._request(
                "initialize",
                {
                    "protocolVersion": self.PROTOCOL_VERSION,
                    "capabilities": {
                        "roots": {"listChanged": True},
                        "sampling": {},
                    },
                    "clientInfo": {
                        "name": "familiar",
                        "version": "1.5.0",
                    },
                },
            )

            self._server_info = response.get("serverInfo", {})
            self._capabilities = response.get("capabilities", {})

            logger.info(
                f"Connected to MCP server '{self.name}': "
                f"{self._server_info.get('name', 'unknown')} "
                f"v{self._server_info.get('version', '?')}"
            )

            # Send initialized notification
            await self._notify("notifications/initialized", {})

            self._initialized = True

        except Exception as e:
            await self._transport.disconnect()
            raise MCPConnectionError(f"Failed to initialize MCP protocol: {e}")

    async def disconnect(self) -> None:
        """Disconnect from the MCP server."""
        self._initialized = False
        self._tools = None
        self._resources = None
        self._prompts = None
        await self._transport.disconnect()

    async def reconnect(self) -> None:
        """Reconnect to the server."""
        await self.disconnect()
        await self.connect()

    # ----------------------------------------------------------
    # Tools
    # ----------------------------------------------------------

    async def list_tools(self, force_refresh: bool = False) -> List[MCPTool]:
        """
        List available tools from the server.

        Args:
            force_refresh: If True, fetch fresh list from server

        Returns:
            List of MCPTool objects
        """
        if self._tools is not None and not force_refresh:
            return self._tools

        response = await self._request("tools/list", {})

        self._tools = [
            MCPTool(
                name=t["name"],
                description=t.get("description", ""),
                input_schema=t.get("inputSchema", {}),
                annotations=t.get("annotations", {}),
            )
            for t in response.get("tools", [])
        ]

        logger.info(f"Server '{self.name}' has {len(self._tools)} tools")
        return self._tools

    async def call_tool(
        self,
        name: str,
        arguments: Optional[Dict[str, Any]] = None,
    ) -> MCPToolResult:
        """
        Call a tool on the server.

        Args:
            name: Tool name
            arguments: Tool arguments (optional)

        Returns:
            MCPToolResult with content blocks
        """
        # Check if tool is allowed
        if not self.config.is_tool_allowed(name):
            raise MCPToolError(
                f"Tool '{name}' is blocked by configuration",
                tool_name=name,
            )

        response = await self._request(
            "tools/call",
            {
                "name": name,
                "arguments": arguments or {},
            },
        )

        return MCPToolResult(
            content=response.get("content", []),
            is_error=response.get("isError", False),
        )

    # ----------------------------------------------------------
    # Resources
    # ----------------------------------------------------------

    async def list_resources(self, force_refresh: bool = False) -> List[MCPResource]:
        """List available resources from the server."""
        if self._resources is not None and not force_refresh:
            return self._resources

        response = await self._request("resources/list", {})

        self._resources = [
            MCPResource(
                uri=r["uri"],
                name=r["name"],
                description=r.get("description"),
                mime_type=r.get("mimeType"),
            )
            for r in response.get("resources", [])
        ]

        return self._resources

    async def read_resource(self, uri: str) -> Dict[str, Any]:
        """Read a resource by URI."""
        response = await self._request("resources/read", {"uri": uri})
        return response.get("contents", [])

    # ----------------------------------------------------------
    # Prompts
    # ----------------------------------------------------------

    async def list_prompts(self, force_refresh: bool = False) -> List[MCPPrompt]:
        """List available prompts from the server."""
        if self._prompts is not None and not force_refresh:
            return self._prompts

        response = await self._request("prompts/list", {})

        self._prompts = [
            MCPPrompt(
                name=p["name"],
                description=p.get("description"),
                arguments=p.get("arguments", []),
            )
            for p in response.get("prompts", [])
        ]

        return self._prompts

    async def get_prompt(
        self,
        name: str,
        arguments: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Get a prompt by name with optional arguments."""
        response = await self._request(
            "prompts/get",
            {
                "name": name,
                "arguments": arguments or {},
            },
        )
        return response

    # ----------------------------------------------------------
    # Internal Methods
    # ----------------------------------------------------------

    def _next_id(self) -> int:
        """Generate next request ID."""
        self._request_id += 1
        return self._request_id

    async def _request(
        self,
        method: str,
        params: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Send a JSON-RPC request and wait for response."""
        request_id = self._next_id()

        message = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
            "params": params,
        }

        if isinstance(self._transport, HTTPTransport):
            # HTTP is synchronous request/response
            return await self._transport.request(message)

        # Stdio uses async message handling
        future = asyncio.get_event_loop().create_future()
        self._pending_requests[request_id] = future

        try:
            await self._transport.send(message)
            response = await asyncio.wait_for(future, timeout=self.config.timeout)

            if "error" in response:
                error = response["error"]
                raise MCPProtocolError(
                    f"MCP error {error.get('code', '?')}: {error.get('message', 'Unknown error')}"
                )

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

        except asyncio.TimeoutError:
            raise MCPTimeoutError(f"Timeout waiting for response to {method}")
        finally:
            self._pending_requests.pop(request_id, None)

    async def _notify(self, method: str, params: Dict[str, Any]) -> None:
        """Send a JSON-RPC notification (no response expected)."""
        message = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }
        await self._transport.send(message)

    async def _message_handler(self) -> None:
        """Handle incoming messages for stdio transport."""
        while self._transport.is_connected:
            try:
                message = await self._transport.receive()

                # Check if it's a response to a pending request
                if "id" in message and message["id"] in self._pending_requests:
                    future = self._pending_requests[message["id"]]
                    if not future.done():
                        future.set_result(message)

                # Handle notifications
                elif "method" in message and "id" not in message:
                    await self._handle_notification(message)

            except MCPTimeoutError:
                continue
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error handling MCP message: {e}")

    async def _handle_notification(self, message: Dict[str, Any]) -> None:
        """Handle server notifications."""
        method = message.get("method", "")
        _params = message.get("params", {})  # available for notification payload (future use)

        if method == "notifications/tools/list_changed":
            # Invalidate tools cache
            self._tools = None
            logger.info(f"Server '{self.name}' tools changed")

        elif method == "notifications/resources/list_changed":
            # Invalidate resources cache
            self._resources = None
            logger.info(f"Server '{self.name}' resources changed")

        elif method == "notifications/prompts/list_changed":
            # Invalidate prompts cache
            self._prompts = None
            logger.info(f"Server '{self.name}' prompts changed")

        else:
            logger.debug(f"Unhandled notification: {method}")
