"""MCP client for the EasyClaw knowledge base server.

JSON-RPC 2.0 over Streamable HTTP transport, with SSE fallback.
Synchronous — runs in the agent's background thread.

Transport auto-detection:
  1. Try Streamable HTTP (POST /mcp) — simpler, FastMCP >= 2.0
  2. Fall back to SSE (GET /sse + POST to endpoint) — legacy FastMCP
"""

from __future__ import annotations

import json
import logging
import queue
import threading
from typing import Any

import httpx

from ..constants import MCP_HTTP_TIMEOUT

logger = logging.getLogger(__name__)


class MCPError(Exception):
    """Error from the MCP server or transport layer."""


class MCPClient:
    """JSON-RPC 2.0 client for a remote MCP server.

    Synchronous — all methods block. Safe to call from the agent's
    background thread only.
    """

    def __init__(self, base_url: str, api_key: str = "") -> None:
        self.base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._session_id: str | None = None
        self._request_counter = 0
        self._initialized = False
        self._transport: str | None = None  # "streamable" or "sse"

        # HTTP client for JSON-RPC requests (both transports)
        self._http = httpx.Client(timeout=MCP_HTTP_TIMEOUT)

        # SSE transport state
        self._sse_client: httpx.Client | None = None
        self._sse_thread: threading.Thread | None = None
        self._sse_endpoint: str | None = None
        self._sse_ready = threading.Event()
        self._sse_responses: dict[str, queue.Queue] = {}
        self._sse_lock = threading.Lock()
        self._sse_running = False

    def _next_id(self) -> str:
        self._request_counter += 1
        return str(self._request_counter)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def connect(self) -> bool:
        """Initialize the MCP connection. Returns True on success."""
        # Try Streamable HTTP first
        try:
            if self._try_streamable():
                self._transport = "streamable"
                logger.info("MCP: connected via Streamable HTTP")
                return True
        except Exception as exc:
            logger.debug("MCP Streamable HTTP failed: %s", exc)

        # Fall back to SSE
        try:
            if self._try_sse():
                self._transport = "sse"
                logger.info("MCP: connected via SSE transport")
                return True
        except Exception as exc:
            logger.debug("MCP SSE failed: %s", exc)

        logger.error("MCP: could not connect via any transport")
        return False

    def list_tools(self) -> list[dict]:
        """Get available tools from the MCP server."""
        result = self._rpc("tools/list")
        return result.get("tools", [])

    def call_tool(self, name: str, arguments: dict) -> str:
        """Call an MCP tool and return the text result."""
        result = self._rpc("tools/call", {"name": name, "arguments": arguments})

        # MCP returns content as array of content blocks
        content_blocks = result.get("content", [])
        texts = []
        for block in content_blocks:
            if block.get("type") == "text":
                texts.append(block.get("text", ""))
        return "\n".join(texts) if texts else json.dumps(result)

    def close(self) -> None:
        """Release all resources."""
        self._sse_running = False
        self._http.close()
        if self._sse_client:
            self._sse_client.close()

    # ------------------------------------------------------------------
    # Unified RPC dispatcher
    # ------------------------------------------------------------------

    def _rpc(self, method: str, params: dict | None = None) -> dict:
        """Route a JSON-RPC call to the active transport."""
        if self._transport == "streamable":
            return self._streamable_rpc(method, params)
        elif self._transport == "sse":
            return self._sse_rpc(method, params)
        raise MCPError("Not connected")

    # ------------------------------------------------------------------
    # Transport: Streamable HTTP
    # ------------------------------------------------------------------

    def _try_streamable(self) -> bool:
        """Attempt connection via Streamable HTTP (POST /mcp)."""
        result = self._streamable_rpc("initialize", {
            "protocolVersion": "2025-03-26",
            "capabilities": {},
            "clientInfo": {"name": "easyclaw-installer", "version": "0.4.0"},
        })
        if result and "protocolVersion" in result:
            self._streamable_notify("notifications/initialized")
            self._initialized = True
            return True
        return False

    def _streamable_rpc(
        self, method: str, params: dict | None = None
    ) -> dict:
        """Send a JSON-RPC request via Streamable HTTP, return result."""
        request_id = self._next_id()
        payload: dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
            "id": request_id,
        }
        if params is not None:
            payload["params"] = params

        headers: dict[str, str] = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        if self._session_id:
            headers["Mcp-Session-Id"] = self._session_id

        resp = self._http.post(
            f"{self.base_url}/mcp", json=payload, headers=headers
        )
        resp.raise_for_status()

        # Capture session ID
        sid = resp.headers.get("mcp-session-id")
        if sid:
            self._session_id = sid

        # Handle SSE response (text/event-stream)
        content_type = resp.headers.get("content-type", "")
        if "text/event-stream" in content_type:
            return self._parse_sse_body(resp.text, request_id)

        data = resp.json()
        if "error" in data:
            raise MCPError(data["error"].get("message", "Unknown MCP error"))
        return data.get("result", {})

    def _streamable_notify(
        self, method: str, params: dict | None = None
    ) -> None:
        """Send a JSON-RPC notification (no id, no response)."""
        payload: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            payload["params"] = params

        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        if self._session_id:
            headers["Mcp-Session-Id"] = self._session_id

        self._http.post(
            f"{self.base_url}/mcp", json=payload, headers=headers
        )

    @staticmethod
    def _parse_sse_body(body: str, request_id: str) -> dict:
        """Parse a non-streaming SSE body for a matching result."""
        for line in body.splitlines():
            if line.startswith("data: "):
                try:
                    msg = json.loads(line[6:])
                    if str(msg.get("id", "")) == request_id:
                        if "error" in msg:
                            raise MCPError(
                                msg["error"].get("message", "MCP error")
                            )
                        return msg.get("result", {})
                except json.JSONDecodeError:
                    continue
        return {}

    # ------------------------------------------------------------------
    # Transport: SSE (legacy)
    # ------------------------------------------------------------------

    def _try_sse(self) -> bool:
        """Attempt connection via SSE transport (GET /sse)."""
        self._sse_client = httpx.Client(timeout=None)
        self._sse_running = True

        self._sse_thread = threading.Thread(
            target=self._sse_reader_loop, daemon=True
        )
        self._sse_thread.start()

        # Wait for the endpoint event
        if not self._sse_ready.wait(timeout=10):
            self._sse_running = False
            return False

        result = self._sse_rpc("initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "easyclaw-installer", "version": "0.4.0"},
        })
        if result and "protocolVersion" in result:
            self._sse_notify("notifications/initialized")
            self._initialized = True
            return True
        return False

    def _sse_reader_loop(self) -> None:
        """Background thread: read SSE events, dispatch responses."""
        try:
            with self._sse_client.stream(
                "GET", f"{self.base_url}/sse"
            ) as resp:
                resp.raise_for_status()
                event_type = ""
                for line in resp.iter_lines():
                    if not self._sse_running:
                        return

                    if line.startswith("event: "):
                        event_type = line[7:].strip()
                    elif line.startswith("data: "):
                        data_str = line[6:]

                        if event_type == "endpoint":
                            # Resolve relative endpoint URL
                            endpoint = data_str.strip()
                            if not endpoint.startswith("http"):
                                endpoint = self.base_url + endpoint
                            self._sse_endpoint = endpoint
                            self._sse_ready.set()

                        elif event_type == "message":
                            try:
                                msg = json.loads(data_str)
                            except json.JSONDecodeError:
                                continue
                            msg_id = str(msg.get("id", ""))
                            with self._sse_lock:
                                q = self._sse_responses.get(msg_id)
                            if q is not None:
                                q.put(msg)

                        event_type = ""
                    elif not line:
                        event_type = ""

        except Exception as exc:
            logger.debug("MCP SSE reader stopped: %s", exc)
            self._sse_running = False

    def _sse_rpc(self, method: str, params: dict | None = None) -> dict:
        """Send JSON-RPC via SSE transport, wait for response."""
        request_id = self._next_id()
        payload: dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
            "id": request_id,
        }
        if params is not None:
            payload["params"] = params

        # Register response queue
        resp_queue: queue.Queue = queue.Queue()
        with self._sse_lock:
            self._sse_responses[request_id] = resp_queue

        try:
            self._http.post(
                self._sse_endpoint,
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            msg = resp_queue.get(timeout=MCP_HTTP_TIMEOUT)
            if "error" in msg:
                raise MCPError(msg["error"].get("message", "MCP error"))
            return msg.get("result", {})
        finally:
            with self._sse_lock:
                self._sse_responses.pop(request_id, None)

    def _sse_notify(
        self, method: str, params: dict | None = None
    ) -> None:
        """Send a JSON-RPC notification via SSE transport."""
        payload: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            payload["params"] = params
        self._http.post(
            self._sse_endpoint,
            json=payload,
            headers={"Content-Type": "application/json"},
        )
