"""STDIO transport — spawns an MCP server as a subprocess and communicates via stdin/stdout."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shlex
import sys
from typing import Any

logger = logging.getLogger("mcp_forge.transport.stdio")


class StdioTransport:
    """Manages a subprocess MCP server communicating over STDIO with JSON-RPC 2.0.

    This transport spawns the server command as a child process, sends JSON-RPC
    messages to its stdin, and reads responses from its stdout. Stderr is captured
    separately for diagnostics.
    """

    def __init__(
        self,
        command: str,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        timeout: float = 10.0,
    ) -> None:
        self.command = command
        self.env = env
        self.cwd = cwd
        self.timeout = timeout
        self._process: asyncio.subprocess.Process | None = None
        self._request_id = 0
        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}
        self._reader_task: asyncio.Task[None] | None = None
        self._stderr_lines: list[str] = []
        self._initialized = False

    async def start(self) -> None:
        """Start the server subprocess and begin reading responses."""
        merged_env = {**os.environ}
        if self.env:
            merged_env.update(self.env)

        if sys.platform == "win32":
            # On Windows, shlex.split mangles backslash paths (C:\Users -> C:Users).
            # Use posix=False to preserve them.
            args = shlex.split(self.command, posix=False)
        else:
            args = shlex.split(self.command)

        self._process = await asyncio.create_subprocess_exec(
            *args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=merged_env,
            cwd=self.cwd,
        )
        self._reader_task = asyncio.create_task(self._read_loop())
        self._stderr_task = asyncio.create_task(self._read_stderr())

    async def _read_loop(self) -> None:
        """Continuously read JSON-RPC responses from the server's stdout."""
        assert self._process is not None
        assert self._process.stdout is not None

        try:
            while True:
                line = await self._process.stdout.readline()
                if not line:
                    break

                line_str = line.decode("utf-8").strip()
                if not line_str:
                    continue

                try:
                    msg = json.loads(line_str)
                except json.JSONDecodeError:
                    logger.debug("Non-JSON line from server: %s", line_str)
                    continue

                # Handle JSON-RPC response (has "id")
                msg_id = msg.get("id")
                if msg_id is not None and msg_id in self._pending:
                    future = self._pending.pop(msg_id)
                    if not future.done():
                        future.set_result(msg)
                # Handle JSON-RPC notification (no "id") — log it
                elif msg.get("method"):
                    logger.debug("Server notification: %s", msg.get("method"))

        except asyncio.CancelledError:
            pass
        except Exception as exc:
            logger.error("Reader loop error: %s", exc)
            # Fail all pending futures
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(exc)
            self._pending.clear()

    async def _read_stderr(self) -> None:
        """Capture stderr for diagnostics."""
        assert self._process is not None
        assert self._process.stderr is not None

        try:
            while True:
                line = await self._process.stderr.readline()
                if not line:
                    break
                decoded = line.decode("utf-8").rstrip()
                self._stderr_lines.append(decoded)
                logger.debug("Server stderr: %s", decoded)
        except asyncio.CancelledError:
            pass

    async def send_request(
        self,
        method: str,
        params: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Send a JSON-RPC request and wait for the response.

        Args:
            method: The JSON-RPC method name.
            params: Optional parameters for the request.
            timeout: Override the default timeout for this request.

        Returns:
            The JSON-RPC response dict.

        Raises:
            TimeoutError: If the server doesn't respond in time.
            ConnectionError: If the server process is not running.
            RuntimeError: If the server returns a JSON-RPC error.
        """
        if self._process is None or self._process.stdin is None:
            raise ConnectionError("Server process is not running")

        if self._process.returncode is not None:
            stderr_tail = "\n".join(self._stderr_lines[-20:])
            raise ConnectionError(
                f"Server process exited with code {self._process.returncode}.\n"
                f"Last stderr:\n{stderr_tail}"
            )

        self._request_id += 1
        req_id = self._request_id

        request: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": req_id,
            "method": method,
        }
        if params is not None:
            request["params"] = params

        loop = asyncio.get_event_loop()
        future: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[req_id] = future

        msg_bytes = (json.dumps(request) + "\n").encode("utf-8")
        self._process.stdin.write(msg_bytes)
        await self._process.stdin.drain()

        effective_timeout = timeout if timeout is not None else self.timeout

        try:
            response = await asyncio.wait_for(future, timeout=effective_timeout)
        except asyncio.TimeoutError:
            self._pending.pop(req_id, None)
            raise TimeoutError(
                f"Server did not respond to '{method}' within {effective_timeout}s"
            ) from None

        if "error" in response:
            err = response["error"]
            raise RuntimeError(
                f"JSON-RPC error {err.get('code', '?')}: {err.get('message', 'unknown')}"
            )

        return response

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

        notification: dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
        }
        if params is not None:
            notification["params"] = params

        msg_bytes = (json.dumps(notification) + "\n").encode("utf-8")
        self._process.stdin.write(msg_bytes)
        await self._process.stdin.drain()

    async def stop(self) -> None:
        """Terminate the server subprocess and clean up."""
        if self._reader_task and not self._reader_task.done():
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass

        if hasattr(self, "_stderr_task") and self._stderr_task and not self._stderr_task.done():
            self._stderr_task.cancel()
            try:
                await self._stderr_task
            except asyncio.CancelledError:
                pass

        if self._process and self._process.returncode is None:
            try:
                self._process.terminate()
                try:
                    await asyncio.wait_for(self._process.wait(), timeout=5.0)
                except asyncio.TimeoutError:
                    self._process.kill()
                    await self._process.wait()
            except ProcessLookupError:
                pass

        # Fail remaining pending futures
        for future in self._pending.values():
            if not future.done():
                future.set_exception(ConnectionError("Transport stopped"))
        self._pending.clear()

    @property
    def stderr_output(self) -> str:
        """Get all captured stderr output from the server."""
        return "\n".join(self._stderr_lines)

    @property
    def is_alive(self) -> bool:
        """Check if the server process is still running."""
        return self._process is not None and self._process.returncode is None

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

    async def __aexit__(self, *args: Any) -> None:
        await self.stop()
