"""
MCP Client module.

Provides a client for the Model Context Protocol (MCP) to interact with MCP servers.
"""

from __future__ import annotations

import os
from contextlib import AsyncExitStack
from collections.abc import Mapping, Sequence
from typing import Optional, Any, List

from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.types import Tool, Resource, Prompt


class GenericMcpClient:
    """
    A generic client for connecting to MCP Servers via SSE or Stdio.
    Wraps the official python SDK to provide a unified, simple interface.
    """

    def __init__(
        self,
        url: Optional[str] = None,
        command: Optional[str] = None,
        args: Optional[Sequence[str]] = None,
        env: Optional[Mapping[str, str]] = None,
    ):
        """
        Initialize the client.

        Args:
            url (str): The URL for SSE connection (e.g., "http://localhost:8000/sse").
            command (str): The command to run for Stdio connection (e.g., "uv").
            args (Sequence[str]): Arguments for the command (e.g., ["run", "python", "script.py"]).
            env (Mapping[str, str]): Environment variables for the Stdio process.
        """
        self.url = url
        self.command = command
        self.args = list(args) if args else []
        self.env = env

        if not self.url and not self.command:
            raise ValueError(
                "Either 'url' (SSE) or 'command' (Stdio) must be provided."
            )

        self._exit_stack = AsyncExitStack()
        self._session: Optional[ClientSession] = None

    async def __aenter__(self) -> "GenericMcpClient":
        """
        Establish connection and initialize session.
        """
        try:
            if self.url:
                # SSE Connection
                read, write = await self._exit_stack.enter_async_context(
                    sse_client(self.url)
                )
            else:
                # Stdio Connection
                server_params = StdioServerParameters(
                    command=self.command,
                    args=self.args,
                    env={**os.environ, **(self.env or {})},
                )
                read, write = await self._exit_stack.enter_async_context(
                    stdio_client(server_params)
                )

            # Create and initialize session
            self._session = await self._exit_stack.enter_async_context(
                ClientSession(read, write)
            )
            if self.session:
                await self.session.initialize()
            return self

        except Exception as e:
            # Clean up on failure
            await self._exit_stack.aclose()
            raise ConnectionError(f"Failed to connect to MCP server: {e}") from e

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        Close session and transport.
        """
        await self._exit_stack.aclose()

    @property
    def session(self) -> ClientSession:
        if not self._session:
            raise RuntimeError("Client not connected. Use 'async with client:' block.")
        return self._session

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

    async def list_tools(self) -> List[Tool]:
        """List available tools."""
        result = await self.session.list_tools()
        return result.tools

    async def call_tool(
        self, name: str, arguments: Optional[Mapping[str, Any]] = None
    ) -> Any:
        """Call a tool and return the text content of the result."""
        result = await self.session.call_tool(name, arguments=arguments or {})

        # Simple text extraction.
        # For complex results (images, etc.), user might want raw result,
        # but for a "generic client" helper, text is the 90% case.
        texts = [c.text for c in result.content if c.type == "text"]
        return "\n".join(texts)

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

    async def list_resources(self) -> List[Resource]:
        """List available resources."""
        result = await self.session.list_resources()
        return result.resources

    async def read_resource(self, uri: str) -> str:
        """Read a resource and return its text content."""
        result = await self.session.read_resource(uri)
        texts = [c.text for c in result.content if c.type == "text"]
        return "\n".join(texts)

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

    async def list_prompts(self) -> List[Prompt]:
        """List available prompts."""
        result = await self.session.list_prompts()
        return result.prompts

    async def get_prompt(
        self, name: str, arguments: Optional[Mapping[str, Any]] = None
    ) -> Any:
        """Get a prompt."""
        return await self.session.get_prompt(name, arguments=arguments)
