"""
MCP Client - Client for communicating with MCP servers

Handles the MCP protocol lifecycle:
- Connection initialization
- Tool discovery and execution
- Resource management
- Prompt retrieval
"""

import asyncio
from typing import Any, Optional

from groknroll.mcp.transport import MCPTransport, create_transport
from groknroll.mcp.types import (
    MCPCapability,
    MCPMessage,
    MCPMethod,
    MCPPrompt,
    MCPResource,
    MCPServerConfig,
    MCPServerInfo,
    MCPTool,
    MCPToolParameter,
    MCPToolResult,
)
from typing import Callable


class MCPClient:
    """
    MCP Protocol Client

    Manages communication with an MCP server, handling:
    - Protocol initialization and capability negotiation
    - Tool listing and invocation
    - Resource listing and reading
    - Prompt listing and retrieval

    Example:
        config = MCPServerConfig(
            name="filesystem",
            command="npx",
            args=["-y", "@modelcontextprotocol/server-filesystem", "/path"]
        )
        client = MCPClient(config)
        await client.connect()
        tools = await client.list_tools()
        result = await client.call_tool("read_file", path="/etc/hosts")
        await client.disconnect()
    """

    # Client identification
    CLIENT_NAME = "groknroll"
    CLIENT_VERSION = "2.1.3"
    PROTOCOL_VERSION = "2024-11-05"

    def __init__(self, config: MCPServerConfig):
        """
        Initialize MCP client

        Args:
            config: MCP server configuration
        """
        self.config = config
        self._transport: Optional[MCPTransport] = None
        self._server_info: Optional[MCPServerInfo] = None
        self._request_id = 0
        self._pending_requests: dict[int, asyncio.Future] = {}
        self._receive_task: Optional[asyncio.Task] = None

        # Notification callbacks (OpenCode compatible)
        self._on_tools_changed: Optional[Callable] = None
        self._on_resources_changed: Optional[Callable] = None
        self._on_prompts_changed: Optional[Callable] = None

    async def connect(self) -> MCPServerInfo:
        """
        Connect to MCP server and initialize protocol

        Returns:
            Server information including capabilities

        Raises:
            RuntimeError: If connection fails
        """
        # Create transport
        self._transport = create_transport(
            transport_type=self.config.transport.value,
            command=self.config.command,
            args=self.config.args,
            env=self.config.env,
            url=self.config.url,
        )

        # Connect transport
        await self._transport.connect()

        # Start receive loop
        self._receive_task = asyncio.create_task(self._receive_loop())

        # Initialize protocol
        self._server_info = await self._initialize()

        return self._server_info

    async def disconnect(self) -> None:
        """Disconnect from MCP server"""
        # Cancel receive loop
        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
            self._receive_task = None

        # Disconnect transport
        if self._transport:
            await self._transport.disconnect()
            self._transport = None

        self._server_info = None

    async def _receive_loop(self) -> None:
        """Background task to receive and route messages"""
        while self._transport and self._transport.is_connected():
            try:
                message = await self._transport.receive()

                # Handle response to pending request
                if message.id is not None and message.id in self._pending_requests:
                    future = self._pending_requests.pop(message.id)
                    if message.error:
                        future.set_exception(
                            RuntimeError(
                                f"MCP error: {message.error.get('message', 'Unknown error')}"
                            )
                        )
                    else:
                        future.set_result(message.result)

                # Handle notifications (no id)
                elif message.method:
                    await self._handle_notification(message)

            except EOFError:
                break
            except asyncio.CancelledError:
                break
            except Exception:
                # Log error but continue
                pass

    async def _handle_notification(self, message: MCPMessage) -> None:
        """
        Handle incoming notifications from server.

        OpenCode compatible: Handles tool/resource/prompt list change notifications.
        """
        method = message.method

        # Handle tools/list_changed notification (OpenCode compatible)
        if method == MCPMethod.TOOLS_LIST_CHANGED:
            if self._on_tools_changed:
                try:
                    # Refresh tool list
                    tools = await self.list_tools()
                    if self._server_info:
                        self._server_info.tools = tools
                    # Call callback
                    if asyncio.iscoroutinefunction(self._on_tools_changed):
                        await self._on_tools_changed(tools)
                    else:
                        self._on_tools_changed(tools)
                except Exception:
                    pass  # Don't let callback errors break the receive loop

        # Handle resources/list_changed notification
        elif method == MCPMethod.RESOURCES_LIST_CHANGED:
            if self._on_resources_changed:
                try:
                    resources = await self.list_resources()
                    if self._server_info:
                        self._server_info.resources = resources
                    if asyncio.iscoroutinefunction(self._on_resources_changed):
                        await self._on_resources_changed(resources)
                    else:
                        self._on_resources_changed(resources)
                except Exception:
                    pass

        # Handle prompts/list_changed notification
        elif method == MCPMethod.PROMPTS_LIST_CHANGED:
            if self._on_prompts_changed:
                try:
                    prompts = await self.list_prompts()
                    if self._server_info:
                        self._server_info.prompts = prompts
                    if asyncio.iscoroutinefunction(self._on_prompts_changed):
                        await self._on_prompts_changed(prompts)
                    else:
                        self._on_prompts_changed(prompts)
                except Exception:
                    pass

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

        Args:
            method: MCP method name
            params: Request parameters

        Returns:
            Response result
        """
        if not self._transport:
            raise RuntimeError("Not connected")

        # Generate request ID
        self._request_id += 1
        request_id = self._request_id

        # Create future for response
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._pending_requests[request_id] = future

        # Send request
        message = MCPMessage(
            method=method,
            params=params or {},
            id=request_id,
        )
        await self._transport.send(message)

        # Wait for response
        try:
            result = await asyncio.wait_for(future, timeout=30.0)
            return result
        except asyncio.TimeoutError:
            self._pending_requests.pop(request_id, None)
            raise RuntimeError(f"Request timed out: {method}")

    async def _send_notification(
        self, method: str, params: Optional[dict[str, Any]] = None
    ) -> None:
        """
        Send a notification (no response expected)

        Args:
            method: MCP method name
            params: Notification parameters
        """
        if not self._transport:
            raise RuntimeError("Not connected")

        message = MCPMessage(
            method=method,
            params=params or {},
        )
        await self._transport.send(message)

    async def _initialize(self) -> MCPServerInfo:
        """
        Initialize MCP protocol with server

        Returns:
            Server information and capabilities
        """
        # Send initialize request
        result = await self._send_request(
            MCPMethod.INITIALIZE,
            {
                "protocolVersion": self.PROTOCOL_VERSION,
                "capabilities": {
                    "tools": {},
                    "resources": {},
                    "prompts": {},
                },
                "clientInfo": {
                    "name": self.CLIENT_NAME,
                    "version": self.CLIENT_VERSION,
                },
            },
        )

        # Parse server capabilities
        server_info = result.get("serverInfo", {})
        capabilities = result.get("capabilities", {})

        info = MCPServerInfo(
            name=server_info.get("name", self.config.name),
            version=server_info.get("version", "unknown"),
        )

        # Parse capabilities
        if "tools" in capabilities:
            info.capabilities.append(MCPCapability.TOOLS)
        if "resources" in capabilities:
            info.capabilities.append(MCPCapability.RESOURCES)
        if "prompts" in capabilities:
            info.capabilities.append(MCPCapability.PROMPTS)
        if "logging" in capabilities:
            info.capabilities.append(MCPCapability.LOGGING)
        if "sampling" in capabilities:
            info.capabilities.append(MCPCapability.SAMPLING)

        # Send initialized notification
        await self._send_notification(MCPMethod.INITIALIZED)

        return info

    async def list_tools(self) -> list[MCPTool]:
        """
        List tools available from the server

        Returns:
            List of MCP tools
        """
        result = await self._send_request(MCPMethod.TOOLS_LIST)

        tools = []
        for tool_data in result.get("tools", []):
            params = []
            schema = tool_data.get("inputSchema", {})

            # Parse parameters from JSON Schema
            properties = schema.get("properties", {})
            required = schema.get("required", [])

            for param_name, param_schema in properties.items():
                params.append(
                    MCPToolParameter(
                        name=param_name,
                        description=param_schema.get("description", ""),
                        type=param_schema.get("type", "string"),
                        required=param_name in required,
                        enum=param_schema.get("enum"),
                        default=param_schema.get("default"),
                    )
                )

            tools.append(
                MCPTool(
                    name=tool_data.get("name", ""),
                    description=tool_data.get("description", ""),
                    parameters=params,
                )
            )

        return tools

    async def call_tool(self, name: str, **arguments: Any) -> MCPToolResult:
        """
        Call a tool on the MCP server

        Args:
            name: Tool name
            **arguments: Tool arguments

        Returns:
            Tool execution result
        """
        result = await self._send_request(
            MCPMethod.TOOLS_CALL,
            {
                "name": name,
                "arguments": arguments,
            },
        )

        # Parse result
        content = result.get("content", [])
        is_error = result.get("isError", False)

        # Extract text content
        text_content = []
        for item in content:
            if item.get("type") == "text":
                text_content.append(item.get("text", ""))

        return MCPToolResult(
            content="\n".join(text_content) if text_content else content,
            is_error=is_error,
            error_message=result.get("errorMessage"),
        )

    async def list_resources(self) -> list[MCPResource]:
        """
        List resources available from the server

        Returns:
            List of MCP resources
        """
        result = await self._send_request(MCPMethod.RESOURCES_LIST)

        resources = []
        for res_data in result.get("resources", []):
            resources.append(
                MCPResource(
                    uri=res_data.get("uri", ""),
                    name=res_data.get("name", ""),
                    description=res_data.get("description", ""),
                    mime_type=res_data.get("mimeType"),
                )
            )

        return resources

    async def read_resource(self, uri: str) -> str:
        """
        Read a resource from the server

        Args:
            uri: Resource URI

        Returns:
            Resource content
        """
        result = await self._send_request(
            MCPMethod.RESOURCES_READ,
            {"uri": uri},
        )

        contents = result.get("contents", [])
        if contents:
            # Return text content
            for content in contents:
                if "text" in content:
                    return content["text"]
                if "blob" in content:
                    return content["blob"]

        return ""

    async def list_prompts(self) -> list[MCPPrompt]:
        """
        List prompts available from the server

        Returns:
            List of MCP prompts
        """
        result = await self._send_request(MCPMethod.PROMPTS_LIST)

        prompts = []
        for prompt_data in result.get("prompts", []):
            args = []
            for arg in prompt_data.get("arguments", []):
                args.append(
                    MCPToolParameter(
                        name=arg.get("name", ""),
                        description=arg.get("description", ""),
                        type="string",
                        required=arg.get("required", False),
                    )
                )

            prompts.append(
                MCPPrompt(
                    name=prompt_data.get("name", ""),
                    description=prompt_data.get("description", ""),
                    arguments=args,
                )
            )

        return prompts

    async def get_prompt(self, name: str, **arguments: Any) -> dict[str, Any]:
        """
        Get a prompt from the server

        Args:
            name: Prompt name
            **arguments: Prompt arguments

        Returns:
            Prompt content
        """
        result = await self._send_request(
            MCPMethod.PROMPTS_GET,
            {
                "name": name,
                "arguments": arguments,
            },
        )

        return result

    @property
    def server_info(self) -> Optional[MCPServerInfo]:
        """Get server information"""
        return self._server_info

    @property
    def is_connected(self) -> bool:
        """Check if client is connected"""
        return self._transport is not None and self._transport.is_connected()

    def on_tools_changed(self, callback: Callable) -> None:
        """
        Set callback for tools list changed notification.

        OpenCode compatible: Called when MCP server's tool list changes.

        Args:
            callback: Function to call with updated tool list.
                     Can be sync or async: callback(tools: list[MCPTool])
        """
        self._on_tools_changed = callback

    def on_resources_changed(self, callback: Callable) -> None:
        """
        Set callback for resources list changed notification.

        Args:
            callback: Function to call with updated resource list.
                     Can be sync or async: callback(resources: list[MCPResource])
        """
        self._on_resources_changed = callback

    def on_prompts_changed(self, callback: Callable) -> None:
        """
        Set callback for prompts list changed notification.

        Args:
            callback: Function to call with updated prompt list.
                     Can be sync or async: callback(prompts: list[MCPPrompt])
        """
        self._on_prompts_changed = callback
