"""stdio transport for MCP servers.

In the stdio transport, the client spawns an MCP server subprocess and exchanges
newline-delimited JSON-RPC messages via stdin/stdout.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from asyncio.subprocess import Process
from typing import Any

from mcpbridge.exceptions import MCPSchemaValidationError, MCPProtocolError, TransportDisconnectedError
from mcpbridge.transport.base import BaseTransport

logger = logging.getLogger("mcpbridge")


class StdioTransport(BaseTransport):
    """MCP stdio transport implementation."""

    def __init__(
        self,
        command: str,
        args: list[str] | None = None,
        env: dict[str, str] | None = None,
        headers: dict[str, str] | None = None,  # unused but kept for signature consistency
        timeout: int = 30,
    ) -> None:
        self._command = command
        self._args = args or []
        self._env = env or {}
        self._timeout = timeout

        self._proc: Process | None = None
        self._stdout_task: asyncio.Task[None] | None = None
        self._stderr_task: asyncio.Task[None] | None = None
        self._write_lock = asyncio.Lock()
        self._connected = False

        self._pending: dict[str | int, asyncio.Future[dict[str, Any]]] = {}
        self.notifications_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    @property
    def is_connected(self) -> bool:
        return self._connected

    async def connect(self) -> None:
        """Spawn the MCP server subprocess and begin reading stdout."""
        if self._connected:
            return

        env = dict(**{k: v for k, v in self._env.items()})
        # Inherit process environment by default; allow overrides from provided env.
        for k, v in os.environ.items():
            env.setdefault(k, v)

        self._proc = await asyncio.create_subprocess_exec(
            self._command,
            *self._args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        assert self._proc.stdin is not None
        assert self._proc.stdout is not None

        self._connected = True
        self._stdout_task = asyncio.create_task(self._read_stdout_loop(), name="mcpbridge-stdio-stdout")
        self._stderr_task = asyncio.create_task(self._read_stderr_loop(), name="mcpbridge-stdio-stderr")

    async def close(self) -> None:
        """Terminate subprocess and cancel read tasks."""
        self._connected = False
        if self._stdout_task:
            self._stdout_task.cancel()
        if self._stderr_task:
            self._stderr_task.cancel()

        if self._proc:
            try:
                if self._proc.stdin:
                    self._proc.stdin.close()
            except Exception:
                pass

            try:
                await asyncio.wait_for(self._proc.wait(), timeout=2)
            except Exception:
                try:
                    self._proc.kill()
                except Exception:
                    pass
        self._proc = None

        # Fail any pending futures
        for fut in list(self._pending.values()):
            if not fut.done():
                fut.set_exception(TransportDisconnectedError("Transport disconnected"))
        self._pending.clear()

    async def send_notification(self, method: str, params: dict[str, Any]) -> None:
        """Send JSON-RPC notification."""
        if not self._connected or not self._proc or not self._proc.stdin:
            raise TransportDisconnectedError("Transport not connected")

        payload: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            payload["params"] = params
        await self._write_json(payload)

    async def send(self, message: dict[str, Any]) -> dict[str, Any]:
        """Send JSON-RPC request and wait for matching response."""
        if not self._connected or not self._proc or not self._proc.stdin:
            raise TransportDisconnectedError("Transport not connected")

        msg_id = message.get("id")
        if msg_id is None:
            raise MCPProtocolError("send() requires a JSON-RPC `id` field")
        key = msg_id

        loop = asyncio.get_running_loop()
        fut: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[key] = fut
        try:
            await self._write_json(message)
            resp = await asyncio.wait_for(fut, timeout=self._timeout)
            return resp
        except asyncio.TimeoutError as e:
            if key in self._pending:
                self._pending.pop(key, None)
            raise TimeoutError("Transport timed out") from e
        finally:
            # ensure cleanup
            self._pending.pop(key, None)

    async def _write_json(self, payload: dict[str, Any]) -> None:
        """Write a JSON-RPC payload to server stdin."""
        assert self._proc is not None
        assert self._proc.stdin is not None
        data = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
        async with self._write_lock:
            self._proc.stdin.write((data + "\n").encode("utf-8"))
            await self._proc.stdin.drain()

    async def _read_stdout_loop(self) -> None:
        """Read JSON-RPC messages from stdout and dispatch by id."""
        assert self._proc is not None
        assert self._proc.stdout is not None
        try:
            while True:
                line = await self._proc.stdout.readline()
                if not line:
                    break
                text = line.decode("utf-8", errors="replace").strip()
                if not text:
                    continue
                try:
                    msg = json.loads(text)
                except Exception as e:
                    logger.exception("Malformed JSON from MCP stdio: %r", text)
                    raise MCPProtocolError("Malformed JSON from MCP server") from e

                msg_id = msg.get("id")
                method = msg.get("method")
                # Notifications: no `id`, but has method.
                if msg_id is None and isinstance(method, str) and "result" not in msg and "error" not in msg:
                    await self.notifications_queue.put(msg)
                    continue

                # Responses/requests: resolve pending by id if present.
                if msg_id is not None and msg_id in self._pending:
                    fut = self._pending.get(msg_id)
                    if fut and not fut.done():
                        fut.set_result(msg)
                    continue

                # Otherwise: ignore unknown or unsolicited messages.
        except Exception as e:
            # Fail all pending futures and close.
            for fut in list(self._pending.values()):
                if not fut.done():
                    fut.set_exception(TransportDisconnectedError(f"Transport disconnected: {e}"))
            self._pending.clear()
            self._connected = False
            logger.error("stdio transport reader stopped: %s", e)

    async def _read_stderr_loop(self) -> None:
        """Read stderr for debug logging."""
        if not self._proc or not self._proc.stderr:
            return
        try:
            while True:
                line = await self._proc.stderr.readline()
                if not line:
                    break
                text = line.decode("utf-8", errors="replace").rstrip()
                if text:
                    logger.debug("[mcp-stdio stderr] %s", text)
        except asyncio.CancelledError:
            return

