"""MCP stdio simulator — spawn a server subprocess and communicate via JSON-RPC."""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from typing import Any

from hatchdx import HdxError
from hatchdx.utils.jsonrpc import build_notification, build_request


class SimulatorError(HdxError):
    """Raised when the simulator encounters an unrecoverable problem."""


@dataclass
class ToolInfo:
    """Metadata about a tool exposed by the server."""

    name: str
    description: str
    input_schema: dict[str, Any]
    annotations: dict[str, Any] = field(default_factory=dict)


@dataclass
class ToolCallResult:
    """The result of calling a tool on the server."""

    content: list[dict[str, Any]]
    is_error: bool


class StdioSimulator:
    """Communicate with an MCP server over stdio using JSON-RPC 2.0.

    Usage::

        async with StdioSimulator(["python", "-m", "my_server"]) as sim:
            tools = await sim.list_tools()
            result = await sim.call_tool("my_tool", {"arg": "value"})
    """

    def __init__(
        self,
        command: list[str],
        *,
        timeout: float = 5.0,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
    ) -> None:
        self._command = command
        self._timeout = timeout
        self._env = env
        self._cwd = cwd
        self._process: asyncio.subprocess.Process | None = None
        self._request_id = 0
        self._initialized = False

    # -- Context manager ------------------------------------------------------

    async def __aenter__(self) -> StdioSimulator:
        await self.start()
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.shutdown()

    # -- Lifecycle ------------------------------------------------------------

    async def start(self) -> dict[str, Any]:
        """Spawn the server subprocess and perform the initialization handshake.

        Returns the server's ``initialize`` result (capabilities, serverInfo, etc.).
        """
        try:
            self._process = await asyncio.create_subprocess_exec(
                *self._command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=self._env,
                cwd=self._cwd,
            )
        except FileNotFoundError:
            raise SimulatorError(
                f"Server command not found: {self._command[0]}\n"
                "Make sure the server is installed and accessible in PATH."
            )
        except OSError as e:
            raise SimulatorError(f"Failed to start server process: {e}")

        # Step 1: Send initialize request
        init_result = await self._send_request(
            "initialize",
            {
                "protocolVersion": "2025-03-26",
                "capabilities": {},
                "clientInfo": {
                    "name": "hdx-test",
                    "version": "0.1.0",
                },
            },
        )

        # Step 2: Send initialized notification (no response expected)
        await self._send_notification("notifications/initialized")

        self._initialized = True
        return init_result

    async def shutdown(self) -> None:
        """Cleanly shut down the server subprocess."""
        if self._process is None:
            return

        proc = self._process
        self._process = None
        self._initialized = False

        if proc.returncode is not None:
            # Already exited
            return

        # Close stdin to signal the server to exit
        if proc.stdin is not None:
            proc.stdin.close()
            try:
                await proc.stdin.wait_closed()
            except (BrokenPipeError, ConnectionResetError):
                pass

        try:
            await asyncio.wait_for(proc.wait(), timeout=3.0)
        except asyncio.TimeoutError:
            proc.terminate()
            try:
                await asyncio.wait_for(proc.wait(), timeout=2.0)
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()

    # -- MCP operations -------------------------------------------------------

    async def list_tools(self) -> list[ToolInfo]:
        """Send ``tools/list`` and return the list of available tools."""
        self._ensure_initialized()
        result = await self._send_request("tools/list", {})
        tools_raw = result.get("tools", [])
        return [
            ToolInfo(
                name=t["name"],
                description=t.get("description", ""),
                input_schema=t.get("inputSchema", {}),
                annotations=t.get("annotations", {}),
            )
            for t in tools_raw
        ]

    async def call_tool(
        self, name: str, arguments: dict[str, Any] | None = None
    ) -> ToolCallResult:
        """Send ``tools/call`` and return the result."""
        self._ensure_initialized()
        result = await self._send_request(
            "tools/call",
            {"name": name, "arguments": arguments or {}},
        )
        return ToolCallResult(
            content=result.get("content", []),
            is_error=result.get("isError", False),
        )

    # -- Internal helpers -----------------------------------------------------

    def _ensure_initialized(self) -> None:
        if not self._initialized:
            raise SimulatorError(
                "Simulator not initialized. Call start() or use as async context manager."
            )

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

    async def _send_request(
        self, method: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Send a JSON-RPC request and wait for the matching response."""
        if self._process is None or self._process.stdin is None:
            raise SimulatorError("Server process is not running.")

        request_id = self._next_id()
        line = build_request(method, params, request_id) + "\n"
        self._process.stdin.write(line.encode())
        await self._process.stdin.drain()

        return await self._read_response(request_id)

    async def _send_notification(self, method: str, params: dict[str, Any] | None = None) -> None:
        """Send a JSON-RPC notification (no id, no response expected)."""
        if self._process is None or self._process.stdin is None:
            raise SimulatorError("Server process is not running.")

        line = build_notification(method, params) + "\n"
        self._process.stdin.write(line.encode())
        await self._process.stdin.drain()

    async def _read_response(self, expected_id: int) -> dict[str, Any]:
        """Read lines from stdout until we get a JSON-RPC response with the expected id."""
        if self._process is None or self._process.stdout is None:
            raise SimulatorError("Server process is not running.")

        try:
            while True:
                raw_line = await asyncio.wait_for(
                    self._process.stdout.readline(),
                    timeout=self._timeout,
                )

                if not raw_line:
                    # EOF — server closed stdout
                    stderr_output = ""
                    if self._process.stderr is not None:
                        stderr_bytes = await self._process.stderr.read()
                        stderr_output = stderr_bytes.decode(errors="replace")
                    raise SimulatorError(
                        "Server process exited unexpectedly."
                        + (f"\nStderr: {stderr_output}" if stderr_output else "")
                    )

                line = raw_line.decode().strip()
                if not line:
                    continue

                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    # Skip non-JSON lines (servers might emit debug output)
                    continue

                # Skip notifications from the server (no id field)
                if "id" not in msg:
                    continue

                if msg["id"] != expected_id:
                    # Not our response — skip it (shouldn't happen in normal flow)
                    continue

                # Check for JSON-RPC error
                if "error" in msg:
                    err = msg["error"]
                    raise SimulatorError(
                        f"JSON-RPC error {err.get('code', '?')}: {err.get('message', 'unknown')}"
                    )

                return msg.get("result", {})

        except asyncio.TimeoutError:
            raise SimulatorError(
                f"Timeout waiting for response (>{self._timeout}s). "
                "The server may be hanging or not writing to stdout."
            ) from None
