from __future__ import annotations

"""Async MCP client over stdio or Streamable HTTP transport."""

import asyncio
import json
from dataclasses import asdict, dataclass, field
from itertools import count
from typing import Any

import httpx

from agentrspy.core import McpError, Tool, ToolOutput

from .adapter import McpToolAdapter
from .protocol import McpCallToolResult, McpMessage, McpTool

MCP_PROTOCOL_VERSION = "2025-03-26"


@dataclass(slots=True)
class WebMcpOptions:
    """Configuration for connecting to a web MCP endpoint."""

    headers_map: dict[str, str] = field(default_factory=dict)
    api_key_value: str | None = None
    api_key_header_value: str | None = None
    api_key_prefix_value: str | None = None

    @classmethod
    def new(cls) -> "WebMcpOptions":
        return cls()

    def api_key(self, api_key: str) -> "WebMcpOptions":
        self.api_key_value = api_key
        return self

    def api_key_header(self, header_name: str) -> "WebMcpOptions":
        self.api_key_header_value = header_name
        return self

    def api_key_prefix(self, prefix: str) -> "WebMcpOptions":
        self.api_key_prefix_value = prefix
        return self

    def bearer_auth(self, token: str) -> "WebMcpOptions":
        self.api_key_value = token
        self.api_key_header_value = "Authorization"
        self.api_key_prefix_value = "Bearer "
        return self

    def header(self, name: str, value: str) -> "WebMcpOptions":
        self.headers_map[name] = value
        return self

    def headers(self, headers: dict[str, str]) -> "WebMcpOptions":
        self.headers_map.update(headers)
        return self

    def into_headers(self) -> dict[str, str]:
        headers = dict(self.headers_map)
        if self.api_key_value is not None:
            header_name = self.api_key_header_value or "Authorization"
            if self.api_key_prefix_value is not None:
                header_value = f"{self.api_key_prefix_value}{self.api_key_value}"
            elif header_name.lower() == "authorization":
                header_value = f"Bearer {self.api_key_value}"
            else:
                header_value = self.api_key_value
            headers.setdefault(header_name, header_value)
        return headers


@dataclass(slots=True)
class _StdioTransport:
    process: asyncio.subprocess.Process


@dataclass(slots=True)
class _HttpTransport:
    client: httpx.AsyncClient
    endpoint: str
    headers: dict[str, str]
    session_id: str | None = None


class McpClient:
    """Async MCP client over stdio or Streamable HTTP transport."""

    def __init__(self, transport: _StdioTransport | _HttpTransport) -> None:
        self.transport = transport
        self._ids = count(1)

    @classmethod
    async def spawn(cls, command: str) -> "McpClient":
        parts = command.split()
        if not parts:
            raise McpError.invalid_command()
        try:
            process = await asyncio.create_subprocess_exec(
                parts[0],
                *parts[1:],
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL,
            )
        except OSError as error:
            raise McpError.spawn_failed(str(error)) from error
        client = cls(_StdioTransport(process=process))
        await client.initialize()
        return client

    @classmethod
    async def connect(cls, endpoint: str) -> "McpClient":
        return await cls.connect_with_options(endpoint, WebMcpOptions())

    @classmethod
    async def connect_with_headers(
        cls, endpoint: str, headers: dict[str, str]
    ) -> "McpClient":
        return await cls.connect_with_options(
            endpoint, WebMcpOptions().headers(headers)
        )

    @classmethod
    async def connect_with_api_key(cls, endpoint: str, api_key: str) -> "McpClient":
        return await cls.connect_with_options(
            endpoint, WebMcpOptions().api_key(api_key)
        )

    @classmethod
    async def connect_with_options(
        cls, endpoint: str, options: WebMcpOptions
    ) -> "McpClient":
        client = cls(
            _HttpTransport(
                client=httpx.AsyncClient(),
                endpoint=endpoint,
                headers=options.into_headers(),
            )
        )
        await client.initialize()
        return client

    async def list_tools(self) -> list[McpTool]:
        response = await self.call_method("tools/list", {})
        tools = response.get("tools")
        if tools is None:
            raise McpError.protocol("missing tools field")
        result: list[McpTool] = []
        for item in tools:
            result.append(
                McpTool(
                    name=str(item.get("name") or ""),
                    description=str(item.get("description") or ""),
                    input_schema=item.get("inputSchema")
                    or item.get("input_schema")
                    or {},
                )
            )
        return result

    async def call_tool(self, name: str, arguments: Any) -> ToolOutput:
        response = await self.call_method(
            "tools/call", {"name": name, "arguments": arguments}
        )
        result = McpCallToolResult(
            content=list(response.get("content") or []),
            is_error=bool(response.get("is_error") or False),
        )
        return result.into_tool_output()

    async def initialize(self) -> None:
        await self.call_method(
            "initialize",
            {
                "protocolVersion": MCP_PROTOCOL_VERSION,
                "capabilities": {"tools": {}},
                "clientInfo": {"name": "agentrspy", "version": "0.1.0"},
            },
        )

    async def call_method(self, method: str, params: Any) -> Any:
        message = McpMessage(
            jsonrpc="2.0", id=next(self._ids), method=method, params=params
        )
        if isinstance(self.transport, _StdioTransport):
            response = await _call_stdio_method(self.transport, message)
        else:
            response = await _call_http_method(self.transport, message)
        if response.error is not None:
            raise McpError.response(str(response.error))
        if response.result is None:
            raise McpError.protocol("missing MCP result")
        return response.result


async def spawn_mcp_tools(command: str) -> list[Tool]:
    client = await McpClient.spawn(command)
    tools = await client.list_tools()
    lock = asyncio.Lock()
    return [McpToolAdapter(client=client, definition=tool, lock=lock) for tool in tools]


async def connect_mcp_tools(endpoint: str) -> list[Tool]:
    return await connect_mcp_tools_with_options(endpoint, WebMcpOptions())


async def connect_mcp_tools_with_headers(
    endpoint: str, headers: dict[str, str]
) -> list[Tool]:
    return await connect_mcp_tools_with_options(
        endpoint, WebMcpOptions().headers(headers)
    )


async def connect_mcp_tools_with_api_key(endpoint: str, api_key: str) -> list[Tool]:
    return await connect_mcp_tools_with_options(
        endpoint, WebMcpOptions().api_key(api_key)
    )


async def connect_mcp_tools_with_options(
    endpoint: str, options: WebMcpOptions
) -> list[Tool]:
    client = await McpClient.connect_with_options(endpoint, options)
    tools = await client.list_tools()
    lock = asyncio.Lock()
    return [McpToolAdapter(client=client, definition=tool, lock=lock) for tool in tools]


async def _call_stdio_method(
    transport: _StdioTransport, message: McpMessage
) -> McpMessage:
    if transport.process.stdin is None or transport.process.stdout is None:
        raise McpError.protocol("missing stdio pipes")
    payload = json.dumps(_serialize_message(message)).encode("utf-8") + b"\n"
    transport.process.stdin.write(payload)
    await transport.process.stdin.drain()
    line = await transport.process.stdout.readline()
    if not line.strip():
        raise McpError.protocol("empty MCP response")
    try:
        value = json.loads(line.decode("utf-8"))
    except ValueError as error:
        raise McpError.protocol(str(error)) from error
    return McpMessage(**value)


async def _call_http_method(
    transport: _HttpTransport, message: McpMessage
) -> McpMessage:
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream",
        "MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
        **transport.headers,
    }
    if transport.session_id is not None:
        headers["Mcp-Session-Id"] = transport.session_id
    try:
        response = await transport.client.post(
            transport.endpoint, headers=headers, json=_serialize_message(message)
        )
    except httpx.HTTPError as error:
        raise McpError.protocol(str(error)) from error
    session_id = response.headers.get("Mcp-Session-Id")
    if session_id:
        transport.session_id = session_id
    if response.status_code >= 400:
        raise McpError.response(f"HTTP {response.status_code}: {response.text}")
    content_type = response.headers.get("Content-Type", "")
    if "text/event-stream" in content_type:
        return _parse_sse_response(response.text, message.id)
    try:
        value = response.json()
    except ValueError as error:
        raise McpError.protocol(str(error)) from error
    return McpMessage(**value)


def _parse_sse_response(payload: str, expected_id: int | None) -> McpMessage:
    event_data = ""
    for line in payload.splitlines():
        if line.startswith("data:"):
            event_data += line[5:].lstrip() + "\n"
            continue
        if not line.strip() and event_data.strip():
            message = _extract_response_message(event_data, expected_id)
            if message is not None:
                return message
            event_data = ""
    message = _extract_response_message(event_data, expected_id)
    if message is not None:
        return message
    raise McpError.protocol("missing JSON-RPC response in SSE stream")


def _extract_response_message(
    payload: str, expected_id: int | None
) -> McpMessage | None:
    trimmed = payload.strip()
    if not trimmed or trimmed == "[DONE]":
        return None
    try:
        value = json.loads(trimmed)
    except ValueError as error:
        raise McpError.protocol(str(error)) from error
    if isinstance(value, list):
        for item in value:
            message = McpMessage(**item)
            if expected_id is None or message.id == expected_id:
                return message
        return None
    message = McpMessage(**value)
    if expected_id is None or message.id == expected_id:
        return message
    return None


def _serialize_message(message: McpMessage) -> dict[str, Any]:
    return {key: value for key, value in asdict(message).items() if value is not None}
