"""Stdio transport layer for MCP server communication.

Implements the MCP stdio transport protocol (JSON-RPC 2.0 over stdio)
as specified in the MCP specification version 2024-11-05.
"""

from __future__ import annotations

import json
import logging
import subprocess
import threading
import time
from typing import Any

from cleveragents.mcp.adapter import MCPServerConfig, MCPTransport

logger = logging.getLogger(__name__)

MCP_PROTOCOL_VERSION = "2024-11-05"
DEFAULT_TIMEOUT = 30.0


class StdioMCPTransport(MCPTransport):
    """Concrete stdio transport for MCP server communication.

    Spawns an MCP server as a subprocess and communicates via JSON-RPC 2.0
    messages over stdio using newline-delimited JSON (NDJSON).

    Thread-safety: uses a lock to serialize concurrent calls to the
    subprocess since stdio is not inherently thread-safe.
    """

    def __init__(
        self,
        timeout: float = DEFAULT_TIMEOUT,
    ) -> None:
        self._timeout = timeout
        self._process: subprocess.Popen[bytes] | None = None
        self._lock = threading.RLock()
        self._request_id = 0
        self._connected = False
        self._capabilities: dict[str, Any] = {}

    def connect(self, config: MCPServerConfig) -> dict[str, Any]:
        """Spawn the MCP server subprocess and perform the initialize handshake.

        Parameters
        ----------
        config:
            Server connection configuration with stdio transport settings.

        Returns
        -------
        dict[str, Any]
            Server capabilities from the initialize response.

        Raises
        ------
        ConnectionError
            If the subprocess fails to start or the handshake times out.
        """
        with self._lock:
            if self._connected:
                return self._capabilities

            if config.command is None:
                msg = "stdio transport requires 'command' field"
                raise ConnectionError(msg)

            env = dict(config.env) if config.env else {}
            env["NO_COLOR"] = "1"

            try:
                self._process = subprocess.Popen(
                    [config.command, *config.args],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    env=env,
                    bufsize=0,
                )
            except FileNotFoundError as exc:
                msg = f"Command not found: {config.command}"
                raise ConnectionError(msg) from exc
            except PermissionError as exc:
                msg = f"Permission denied: {config.command}"
                raise ConnectionError(msg) from exc

            self._request_id = 0
            capabilities = self._send_request_sync(
                "initialize",
                {
                    "protocolVersion": MCP_PROTOCOL_VERSION,
                    "capabilities": {"roots": {"listChanged": True}},
                    "clientInfo": {
                        "name": "cleveragents-mcp-adapter",
                        "version": "1.0.0",
                    },
                },
            )

            if capabilities is None:
                self._terminate()
                msg = "Server did not return capabilities during handshake"
                raise ConnectionError(msg)

            self._capabilities = capabilities.get("capabilities", {})
            self._connected = True

            self._send_notification("notifications/initialized", {})

            logger.info("MCP stdio transport connected to '%s'", config.name)
            return capabilities

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        """Send a JSON-RPC method call and return the result.

        Parameters
        ----------
        method:
            The JSON-RPC method name.
        params:
            Method parameters.

        Returns
        -------
        dict[str, Any]
            The JSON-RPC response result.

        Raises
        ------
        TimeoutError
            If the response times out.
        ConnectionError
            If the transport is not connected.
        """
        with self._lock:
            if not self._connected or self._process is None:
                msg = "Transport not connected. Call connect() first."
                raise ConnectionError(msg)

            result = self._send_request_sync(method, params)
            if result is None:
                msg = f"Server did not respond to {method} request"
                raise ConnectionError(msg)
            return result

    def close(self) -> None:
        """Shut down the transport connection and terminate the subprocess."""
        with self._lock:
            if not self._connected:
                return
            self._terminate()
            logger.info("MCP stdio transport closed")

    def _terminate(self) -> None:
        """Terminate the subprocess and reset state."""
        self._connected = False
        self._capabilities = {}
        if self._process is not None:
            try:
                self._process.terminate()
                try:
                    self._process.wait(timeout=5.0)
                except subprocess.TimeoutExpired:
                    self._process.kill()
                    self._process.wait()
            except Exception:
                pass
            self._process = None

    def _send_request_sync(
        self,
        method: str,
        params: dict[str, Any],
    ) -> dict[str, Any] | None:
        """Send a JSON-RPC request synchronously and wait for response.

        Must be called with the lock held.
        """
        if self._process is None:
            return None

        self._request_id += 1
        req_id = self._request_id

        request = {
            "jsonrpc": "2.0",
            "id": req_id,
            "method": method,
            "params": params,
        }

        try:
            if self._process.stdin is None:
                return None
            self._process.stdin.write(json.dumps(request).encode("utf-8") + b"\n")
            self._process.stdin.flush()
        except BrokenPipeError as exc:
            msg = "Subprocess stdin pipe broken"
            raise ConnectionError(msg) from exc

        start = time.monotonic()
        while True:
            if self._process.stdout is None:
                return None
            line = self._process.stdout.readline()
            if not line:
                return None

            if time.monotonic() - start > self._timeout:
                msg = f"Response timeout after {self._timeout}s"
                raise TimeoutError(msg)

            try:
                response = json.loads(line.decode("utf-8"))
            except json.JSONDecodeError:
                continue

            if not isinstance(response, dict):
                continue

            if response.get("id") == req_id:
                if "error" in response:
                    error = response["error"]
                    msg = error.get("message", "Unknown error")
                    raise ConnectionError(f"MCP error: {msg}")
                return response.get("result")

    def _send_notification(self, method: str, params: dict[str, Any]) -> None:
        """Send a JSON-RPC notification (no response expected).

        Must be called with the lock held.
        """
        if self._process is None:
            return

        notification = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }

        if self._process.stdin is None:
            return

        try:
            self._process.stdin.write(json.dumps(notification).encode("utf-8") + b"\n")
            self._process.stdin.flush()
        except BrokenPipeError:
            pass
