from __future__ import annotations

import asyncio
import json
from typing import Any

import httpx
from pydantic import BaseModel

from nexus.protocols.tools.registry import ToolRegistry
from nexus.resilience.errors import MCPError


class MCPTool(BaseModel):
    name: str
    description: str
    input_schema: dict


class MCPClient:
    """MCP client supporting SSE-style and HTTP JSON-RPC transports."""

    def __init__(
        self,
        server_url: str,
        transport: str = "sse",
        headers: dict | None = None,
        timeout: float = 30.0,
    ) -> None:
        self.server_url = server_url.rstrip("/")
        self.transport = transport
        self._http = httpx.AsyncClient(headers=headers or {}, timeout=timeout)
        self._request_id = 1
        self._connected = False

    async def connect(self) -> None:
        """Connect and initialize MCP session."""
        if self.transport == "sse":
            try:
                resp = await self._http.get(f"{self.server_url}/sse")
                if resp.status_code >= 400:
                    raise MCPError(f"Failed SSE connect: {resp.status_code}", server_url=self.server_url)
            except httpx.HTTPError as exc:
                raise MCPError(f"SSE connect error: {exc}", server_url=self.server_url) from exc

        await self._send_rpc(
            method="initialize",
            params={
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "nexus", "version": "1.0"},
            },
            request_id=self._next_request_id(),
        )
        await self._send_rpc(method="notifications/initialized", params={}, request_id=self._next_request_id())
        self._connected = True

    async def list_tools(self) -> list[MCPTool]:
        """List tools from MCP server."""
        result = await self._send_rpc("tools/list", {}, self._next_request_id())
        tools = result.get("result", {}).get("tools", [])
        return [
            MCPTool(
                name=t.get("name", ""),
                description=t.get("description", ""),
                input_schema=t.get("inputSchema", t.get("input_schema", {})),
            )
            for t in tools
        ]

    async def call_tool(self, name: str, arguments: dict) -> Any:
        """Invoke an MCP tool."""
        result = await self._send_rpc(
            "tools/call",
            {"name": name, "arguments": arguments},
            self._next_request_id(),
        )
        content = result.get("result", {}).get("content", [])
        if isinstance(content, list):
            texts = [block.get("text", "") for block in content if isinstance(block, dict)]
            return "\n".join(x for x in texts if x)
        return content

    async def disconnect(self) -> None:
        """Close MCP client resources."""
        self._connected = False
        await self._http.aclose()

    async def _send_rpc(self, method: str, params: dict, request_id: int) -> dict:
        """Send JSON-RPC request via selected transport."""
        payload = {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}
        try:
            if self.transport == "sse":
                resp = await self._http.post(f"{self.server_url}/message", json=payload)
            else:
                resp = await self._http.post(self.server_url, json=payload)
            resp.raise_for_status()
            data = resp.json()
            if "error" in data:
                raise MCPError(f"MCP error: {data['error']}", server_url=self.server_url)
            return data
        except (httpx.HTTPError, ValueError) as exc:
            raise MCPError(f"MCP transport error: {exc}", server_url=self.server_url) from exc

    async def register_to_registry(self, registry: ToolRegistry) -> None:
        """Import MCP tools into local ToolRegistry as async wrappers."""
        tools = await self.list_tools()
        for tool in tools:
            async def _wrapper(_tool_name: str = tool.name, **kwargs: Any) -> Any:
                return await self.call_tool(_tool_name, kwargs)

            _wrapper.__name__ = tool.name
            _wrapper.__doc__ = tool.description
            registry.register(_wrapper, name=tool.name, description=tool.description, tags=["mcp"])

    def _next_request_id(self) -> int:
        rid = self._request_id
        self._request_id += 1
        return rid
