"""
Manages local MCP server connections (stdio and HTTP).

Starts local MCP servers lazily on first request, forwards JSON-RPC
messages, and handles lifecycle (startup, shutdown).
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
import sys
import uuid
from pathlib import Path
from typing import Any, Optional

import httpx

from .config import LocalServerConfig

logger = logging.getLogger(__name__)
MAX_INFLIGHT_PER_SERVER = int(os.environ.get("NEXUS_BRIDGE_MAX_INFLIGHT_PER_SERVER", "16"))
if MAX_INFLIGHT_PER_SERVER <= 0:
    raise ValueError("NEXUS_BRIDGE_MAX_INFLIGHT_PER_SERVER must be > 0.")


class LocalMCPServerConnection:
    """
    Manages a single local MCP server connection (stdio or HTTP).

    For stdio servers: spawns the subprocess, communicates via stdin/stdout.
    For HTTP servers: sends JSON-RPC POST requests.
    """

    def __init__(self, config: LocalServerConfig):
        self.config = config
        self._process: Optional[asyncio.subprocess.Process] = None
        self._http_client: Optional[httpx.AsyncClient] = None
        self._mcp_session_id: Optional[str] = None
        self._session_lock = asyncio.Lock()
        self._started = False
        self._lock = asyncio.Lock()
        self._write_lock = asyncio.Lock()
        self._inflight = 0
        self._inflight_lock = asyncio.Lock()
        self._stdio_reader_task: Optional[asyncio.Task] = None
        self._pending_responses: dict[str, asyncio.Future] = {}
        self._read_buffer = b""

    async def start(self) -> None:
        """Start the local MCP server if not already running."""
        async with self._lock:
            if self._started:
                return

            if self.config.transport == "stdio":
                await self._start_stdio()
                self._stdio_reader_task = asyncio.create_task(self._stdio_reader_loop())
            elif self.config.transport == "http":
                self._http_client = httpx.AsyncClient(
                    timeout=60.0,
                    headers=self.config.headers,
                )
            self._started = True
            logger.info("[ServerManager] Started local server: %s", self.config.name)

    async def _start_stdio(self) -> None:
        """Start a stdio-based MCP server subprocess."""
        env = {**os.environ, **self.config.env}
        cmd = self._resolve_command(self.config.command)
        args = self.config.args

        self._process = await asyncio.create_subprocess_exec(
            cmd, *args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        logger.info(
            "[ServerManager] Spawned stdio process: %s %s (pid=%s)",
            cmd, " ".join(args), self._process.pid,
        )

    @staticmethod
    def _resolve_command(command: Optional[str]) -> str:
        """Resolve a command path or raise a clear error."""
        if not isinstance(command, str) or not command.strip():
            raise ValueError("Server command is required.")

        command = command.strip()
        command_path = Path(command)
        if command_path.is_absolute():
            if not command_path.exists():
                raise ValueError(f"Server command not found: {command}")
            return command

        resolved = shutil.which(command)
        if resolved:
            return resolved

        if sys.platform.startswith("win"):
            for ext in (".cmd", ".bat", ".exe"):
                resolved = shutil.which(f"{command}{ext}")
                if resolved:
                    return resolved

        raise ValueError(
            "Server command not found. Ensure it is installed and on PATH. "
            "On Windows, you may need to set the full path, e.g. "
            r"C:\Program Files\nodejs\npx.cmd"
        )

    async def send_jsonrpc(self, payload: dict) -> dict:
        """
        Send a JSON-RPC message to the local MCP server and return the response.
        """
        if not self._started:
            await self.start()

        async with self._inflight_lock:
            if self._inflight >= MAX_INFLIGHT_PER_SERVER:
                raise RuntimeError(
                    f"Server '{self.config.name}' is overloaded "
                    f"(max in-flight requests: {MAX_INFLIGHT_PER_SERVER})."
                )
            self._inflight += 1

        try:
            if self.config.transport == "stdio":
                return await self._send_stdio(payload)
            return await self._send_http(payload)
        finally:
            async with self._inflight_lock:
                self._inflight -= 1

    async def _send_stdio(self, payload: dict) -> dict:
        """Send JSON-RPC over stdin/stdout using the MCP stdio protocol."""
        if not self._process or self._process.stdin is None or self._process.stdout is None:
            raise RuntimeError(f"Server '{self.config.name}' process is not running.")

        # MCP stdio protocol: JSON-RPC message terminated by newline
        message = json.dumps(payload) + "\n"
        request_id = payload.get("id")
        if request_id is None:
            async with self._write_lock:
                self._process.stdin.write(message.encode("utf-8"))
                await self._process.stdin.drain()
            return {}

        request_id_str = str(request_id)
        response_future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._pending_responses[request_id_str] = response_future

        try:
            async with self._write_lock:
                self._process.stdin.write(message.encode("utf-8"))
                await self._process.stdin.drain()
            response = await asyncio.wait_for(response_future, timeout=60.0)
            if not isinstance(response, dict):
                raise RuntimeError(
                    f"Server '{self.config.name}' returned a non-object JSON-RPC response."
                )
            return response
        finally:
            self._pending_responses.pop(request_id_str, None)

    async def _read_stdout_line(self, *, max_bytes: int) -> bytes:
        """Read a single newline-terminated line with a custom max size."""
        if not self._process or self._process.stdout is None:
            raise RuntimeError(f"Server '{self.config.name}' process is not running.")
        if max_bytes <= 0:
            raise ValueError("max_bytes must be positive.")

        while True:
            newline_index = self._read_buffer.find(b"\n")
            if newline_index != -1:
                line = self._read_buffer[:newline_index + 1]
                self._read_buffer = self._read_buffer[newline_index + 1:]
                return line

            chunk = await self._process.stdout.read(65536)
            if not chunk:
                exit_code = await self._process.wait()
                stderr_text = ""
                if self._process.stderr is not None:
                    stderr_bytes = await self._process.stderr.read()
                    if stderr_bytes:
                        stderr_text = stderr_bytes.decode("utf-8", errors="replace").strip()
                detail = f"exit_code={exit_code}"
                if stderr_text:
                    detail = f"{detail}; stderr={stderr_text}"
                raise RuntimeError(f"Server '{self.config.name}' closed stdout ({detail}).")

            self._read_buffer += chunk
            if len(self._read_buffer) > max_bytes:
                raise RuntimeError(
                    f"Server '{self.config.name}' response exceeded {max_bytes} bytes without newline."
                )

    async def _stdio_reader_loop(self) -> None:
        """Read stdio output continuously and route responses by JSON-RPC id."""
        try:
            while True:
                line = await self._read_stdout_line(max_bytes=5 * 1024 * 1024)
                response = json.loads(line.decode("utf-8"))
                response_id = response.get("id")
                if response_id is None:
                    logger.debug(
                        "[ServerManager] Ignoring stdio message without id from %s: %s",
                        self.config.name,
                        response,
                    )
                    continue
                response_id_str = str(response_id)
                fut = self._pending_responses.get(response_id_str)
                if fut is None:
                    logger.warning(
                        "[ServerManager] Received untracked stdio response id=%s from %s",
                        response_id_str,
                        self.config.name,
                    )
                    continue
                if not fut.done():
                    fut.set_result(response)
        except asyncio.CancelledError:
            pass
        except Exception as exc:
            for fut in self._pending_responses.values():
                if not fut.done():
                    fut.set_exception(exc)
            self._pending_responses.clear()
            logger.error("[ServerManager] stdio reader crashed for %s: %s", self.config.name, exc)

    async def _send_http(self, payload: dict) -> dict:
        """Send JSON-RPC as HTTP POST to the local MCP server."""
        if not self._http_client:
            raise RuntimeError(f"Server '{self.config.name}' HTTP client not initialized.")

        # Check if this is a notification (no id = no response expected)
        is_notification = "id" not in payload
        method = payload.get("method")
        if not isinstance(method, str) or not method.strip():
            raise ValueError("JSON-RPC payload.method is required for HTTP MCP transport.")

        request_headers: dict[str, str] = {}
        # Initialize must start a fresh MCP session negotiation.
        # Do not send a previous session id when re-initializing.
        if method == "initialize":
            async with self._session_lock:
                self._mcp_session_id = None
                return await self._send_http_request(
                    payload=payload,
                    request_headers=request_headers,
                    is_notification=is_notification,
                    method=method,
                )
        if self._mcp_session_id:
            request_headers["Mcp-Session-Id"] = self._mcp_session_id

        return await self._send_http_request(
            payload=payload,
            request_headers=request_headers,
            is_notification=is_notification,
            method=method,
        )

    async def _send_http_request(
        self,
        *,
        payload: dict,
        request_headers: dict[str, str],
        is_notification: bool,
        method: str,
    ) -> dict:
        response = await self._http_client.post(
            self.config.url,
            json=payload,
            headers=request_headers,
        )
        response.raise_for_status()

        response_session_id = response.headers.get("Mcp-Session-Id") or response.headers.get("mcp-session-id")
        if isinstance(response_session_id, str) and response_session_id.strip():
            self._mcp_session_id = response_session_id.strip()

        if is_notification:
            return {}

        response_json = response.json()
        if not isinstance(response_json, dict):
            raise RuntimeError(f"Server '{self.config.name}' returned non-object JSON-RPC payload.")

        if method == "initialize":
            result = response_json.get("result")
            if not isinstance(result, dict):
                raise RuntimeError(
                    f"Server '{self.config.name}' returned invalid initialize result payload."
                )
            result_session_id = result.get("sessionId") or result.get("session_id")
            if isinstance(result_session_id, str) and result_session_id.strip():
                self._mcp_session_id = result_session_id.strip()
            if not isinstance(self._mcp_session_id, str) or not self._mcp_session_id.strip():
                raise RuntimeError(
                    f"Server '{self.config.name}' initialize response missing MCP session id."
                )

        return response_json

    async def stop(self) -> None:
        """Stop the local MCP server."""
        async with self._lock:
            if not self._started:
                return

            if self._stdio_reader_task:
                self._stdio_reader_task.cancel()
                try:
                    await self._stdio_reader_task
                except asyncio.CancelledError:
                    pass
                self._stdio_reader_task = None

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

            if self._http_client:
                await self._http_client.aclose()
                self._http_client = None
            self._mcp_session_id = None
            self._pending_responses.clear()

            self._started = False
            logger.info("[ServerManager] Stopped local server: %s", self.config.name)


class ServerManager:
    """
    Manages all local MCP server connections.

    Servers are started lazily on first request.
    """

    def __init__(self, configs: list[LocalServerConfig]):
        self._configs = {c.name: c for c in configs}
        self._connections: dict[str, LocalMCPServerConnection] = {}

    def get_manifest(self) -> list[dict]:
        """Build the server manifest for the relay registration message."""
        return [
            {"name": c.name, "description": c.description}
            for c in self._configs.values()
        ]

    async def handle_request(self, server_name: str, payload: dict) -> dict:
        """
        Forward a JSON-RPC request to the named local server.

        Starts the server lazily if needed.
        """
        config = self._configs.get(server_name)
        if not config:
            return {
                "jsonrpc": "2.0",
                "id": payload.get("id"),
                "error": {
                    "code": -32601,
                    "message": f"Server '{server_name}' not configured.",
                },
            }

        conn = self._connections.get(server_name)
        if not conn:
            conn = LocalMCPServerConnection(config)
            self._connections[server_name] = conn

        try:
            return await conn.send_jsonrpc(payload)
        except Exception as exc:
            logger.error(
                "[ServerManager] Error forwarding to %s: %s", server_name, exc,
            )
            return {
                "jsonrpc": "2.0",
                "id": payload.get("id"),
                "error": {
                    "code": -32603,
                    "message": f"Local server error: {exc}",
                },
            }

    async def stop_all(self) -> None:
        """Stop all running local MCP servers."""
        for conn in self._connections.values():
            await conn.stop()
        self._connections.clear()
        logger.info("[ServerManager] All local servers stopped.")
