"""MCPTestClient — the core test client for mcp-forge.

Framework-agnostic: works with any MCP server via STDIO, HTTP, or in-memory transport.
"""

from __future__ import annotations

import logging
from typing import Any

from mcp_forge._types import (
    ContentBlock,
    PromptInfo,
    PromptResult,
    ResourceInfo,
    ResourceResult,
    ServerCapabilities,
    ToolInfo,
    ToolResult,
)
from mcp_forge.transport.stdio import StdioTransport

logger = logging.getLogger("mcp_forge.client")

# Protocol version we advertise during initialization
CLIENT_PROTOCOL_VERSION = "2025-06-18"
CLIENT_INFO = {"name": "mcp-forge", "version": "0.1.0"}


class MCPTestClient:
    """Test client for any MCP server.

    Connects to an MCP server, performs the initialization handshake,
    and provides a clean Python API for testing tools, resources, and prompts.

    Usage with STDIO (most common):
        async with MCPTestClient.from_command("python my_server.py") as client:
            tools = await client.list_tools()
            result = await client.call_tool("my_tool", {"param": "value"})

    Usage with FastMCP (in-memory, zero overhead):
        from my_server import mcp
        async with MCPTestClient.from_fastmcp(mcp) as client:
            ...
    """

    def __init__(self, transport: Any) -> None:
        self._transport = transport
        self._capabilities: ServerCapabilities | None = None
        self._server_info: dict[str, Any] = {}
        self._initialized = False

    # ── Factory Methods ──

    @classmethod
    def from_command(
        cls,
        command: str,
        *,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        timeout: float = 10.0,
    ) -> MCPTestClient:
        """Create a client connected to a server launched via STDIO.

        Args:
            command: Shell command to start the server.
                     Examples: "python my_server.py", "npx my-ts-server",
                              "uvx my-mcp-package"
            env: Optional environment variables for the subprocess.
            cwd: Working directory for the subprocess.
            timeout: Connection and request timeout in seconds.

        Returns:
            A connected MCPTestClient instance.
        """
        transport = StdioTransport(command=command, env=env, cwd=cwd, timeout=timeout)
        client = cls(transport)
        return client

    @classmethod
    def from_fastmcp(cls, server: Any) -> MCPTestClient:
        """Create a client connected in-memory to a FastMCP server instance.

        Args:
            server: A FastMCP server instance (the object you decorate tools on).

        Returns:
            A connected MCPTestClient instance.

        Requires:
            pip install 'mcp-forge[fastmcp]'
        """
        from mcp_forge.transport.inmemory import InMemoryTransport

        transport = InMemoryTransport(server)
        client = cls(transport)
        return client

    # ── Initialization ──

    async def _initialize(self) -> None:
        """Perform the MCP initialization handshake."""
        if self._initialized:
            return

        response = await self._transport.send_request(
            "initialize",
            {
                "protocolVersion": CLIENT_PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": CLIENT_INFO,
            },
        )

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

        caps = result.get("capabilities", {})
        self._capabilities = ServerCapabilities(
            tools=bool(caps.get("tools")),
            resources=bool(caps.get("resources")),
            prompts=bool(caps.get("prompts")),
            logging=bool(caps.get("logging")),
            raw=caps,
        )

        # Send initialized notification
        await self._transport.send_notification("notifications/initialized")
        self._initialized = True

        logger.info(
            "Connected to %s (protocol: %s)",
            self._server_info.get("name", "unknown"),
            result.get("protocolVersion", "unknown"),
        )

    # ── Tool Operations ──

    async def list_tools(self) -> list[ToolInfo]:
        """List all tools exposed by the server.

        Returns:
            A list of ToolInfo objects with name, description, and input schema.
        """
        await self._initialize()

        response = await self._transport.send_request("tools/list")
        result = response.get("result", {})
        raw_tools = result.get("tools", [])

        tools = []
        for t in raw_tools:
            tools.append(
                ToolInfo(
                    name=t.get("name", ""),
                    description=t.get("description"),
                    input_schema=t.get("inputSchema", {}),
                    output_schema=t.get("outputSchema"),
                )
            )

        return tools

    async def call_tool(
        self,
        name: str,
        arguments: dict[str, Any] | None = None,
        *,
        timeout: float = 30.0,
    ) -> ToolResult:
        """Call a tool by name and return the result.

        Args:
            name: The tool name to invoke.
            arguments: Tool parameters as a dict.
            timeout: Timeout for this specific call.

        Returns:
            A ToolResult with .is_success, .text, .content, and .raw.
        """
        await self._initialize()

        params: dict[str, Any] = {"name": name}
        if arguments is not None:
            params["arguments"] = arguments

        response = await self._transport.send_request(
            "tools/call", params, timeout=timeout
        )
        result = response.get("result", {})

        content_blocks = []
        for block in result.get("content", []):
            content_blocks.append(
                ContentBlock(
                    type=block.get("type", "text"),
                    text=block.get("text"),
                    data=block.get("data"),
                    mime_type=block.get("mimeType"),
                    raw=block,
                )
            )

        return ToolResult(
            content=content_blocks,
            is_error=result.get("isError", False),
            raw=result,
        )

    async def get_tool(self, name: str) -> ToolInfo | None:
        """Get a specific tool by name, or None if not found."""
        tools = await self.list_tools()
        for t in tools:
            if t.name == name:
                return t
        return None

    async def get_tool_schema(self, name: str) -> dict[str, Any]:
        """Get the input schema for a specific tool.

        Raises:
            KeyError: If the tool doesn't exist.
        """
        tool = await self.get_tool(name)
        if tool is None:
            raise KeyError(f"Tool '{name}' not found on server")
        return tool.input_schema

    async def get_all_schemas(self) -> dict[str, dict[str, Any]]:
        """Get all tool schemas as {name: input_schema}."""
        tools = await self.list_tools()
        return {t.name: t.input_schema for t in tools}

    # ── Resource Operations ──

    async def list_resources(self) -> list[ResourceInfo]:
        """List all resources exposed by the server."""
        await self._initialize()

        response = await self._transport.send_request("resources/list")
        result = response.get("result", {})
        raw_resources = result.get("resources", [])

        return [
            ResourceInfo(
                uri=r.get("uri", ""),
                name=r.get("name"),
                description=r.get("description"),
                mime_type=r.get("mimeType"),
            )
            for r in raw_resources
        ]

    async def read_resource(self, uri: str) -> ResourceResult:
        """Read a resource by URI."""
        await self._initialize()

        response = await self._transport.send_request(
            "resources/read", {"uri": uri}
        )
        result = response.get("result", {})
        contents = result.get("contents", [])

        blocks = []
        for c in contents:
            blocks.append(
                ContentBlock(
                    type="text",
                    text=c.get("text"),
                    data=c.get("blob"),
                    mime_type=c.get("mimeType"),
                    raw=c,
                )
            )

        return ResourceResult(content=blocks)

    # ── Prompt Operations ──

    async def list_prompts(self) -> list[PromptInfo]:
        """List all prompts exposed by the server."""
        await self._initialize()

        response = await self._transport.send_request("prompts/list")
        result = response.get("result", {})
        raw_prompts = result.get("prompts", [])

        return [
            PromptInfo(
                name=p.get("name", ""),
                description=p.get("description"),
                arguments=p.get("arguments", []),
            )
            for p in raw_prompts
        ]

    async def get_prompt(
        self, name: str, arguments: dict[str, Any] | None = None
    ) -> PromptResult:
        """Get a prompt by name with optional arguments."""
        await self._initialize()

        params: dict[str, Any] = {"name": name}
        if arguments:
            params["arguments"] = arguments

        response = await self._transport.send_request("prompts/get", params)
        result = response.get("result", {})

        return PromptResult(
            messages=result.get("messages", []),
            description=result.get("description"),
        )

    # ── Server Info ──

    @property
    def server_name(self) -> str:
        """Server name from initialization."""
        return self._server_info.get("name", "unknown")

    @property
    def server_version(self) -> str:
        """Server version from initialization."""
        return self._server_info.get("version", "unknown")

    @property
    def capabilities(self) -> ServerCapabilities | None:
        """Server capabilities from initialization."""
        return self._capabilities

    # ── Lifecycle ──

    async def connect(self) -> None:
        """Start the transport and initialize the connection."""
        await self._transport.start()
        await self._initialize()

    async def disconnect(self) -> None:
        """Disconnect from the server and stop the transport."""
        await self._transport.stop()
        self._initialized = False

    async def __aenter__(self) -> MCPTestClient:
        await self.connect()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.disconnect()
