"""Minimal newline-delimited JSON-RPC and MCP tools server."""

from __future__ import annotations

import asyncio
from enum import Enum
import io
import json
import logging
from typing import BinaryIO, Final

from .errors import CancelledError, FirefoxMcpError
from .tool_handlers import (
    InvalidArgumentsError,
    ToolHandlers,
    errorResult,
)
from .tool_schemas import TOOLS


MAX_LINE_BYTES: Final = 8 * 1024 * 1024
SUPPORTED_PROTOCOL_VERSIONS: Final = (
    "2025-11-25",
    "2025-06-18",
    "2024-11-05",
)


class ServerState(Enum):
    """Represent the MCP process lifecycle."""

    STARTING = "starting"
    INITIALIZED = "initialized"
    RUNNING = "running"
    STOPPING = "stopping"


class JsonRpcError(Exception):
    """Carry a JSON-RPC error code, message, and optional data."""

    def __init__(
        self,
        code: int,
        message: str,
        data: object | None = None,
    ) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.data = data


class McpServer:
    """Parse MCP requests, dispatch tools, and serialize stdout responses."""

    def __init__(
        self,
        tools: ToolHandlers,
        output: BinaryIO,
        *,
        logger: logging.Logger | None = None,
    ) -> None:
        self.tools = tools
        self.output = output
        self.logger = logger or logging.getLogger("firefox_mcp.mcp")
        self.state = ServerState.STARTING
        self._output_lock = asyncio.Lock()
        self._tasks: dict[object, asyncio.Task[None]] = {}
        self._cancelled: dict[object, asyncio.Event] = {}

    async def run(self, input_stream: BinaryIO) -> None:
        """Read bounded input lines until EOF, then finish active requests."""
        buffer = bytearray()
        discarding = False
        async for chunk in self._readChunks(input_stream):
            for value in chunk:
                if discarding:
                    if value == 0x0a:
                        discarding = False
                    continue
                buffer.append(value)
                if value == 0x0a:
                    await self.processLine(bytes(buffer))
                    buffer.clear()
                    await asyncio.sleep(0)
                elif len(buffer) > MAX_LINE_BYTES:
                    buffer.clear()
                    discarding = True
                    await self._writeError(None, -32700, "Parse error")
        if buffer and not discarding:
            await self.processLine(bytes(buffer))
        if self._tasks:
            await asyncio.gather(
                *tuple(self._tasks.values()),
                return_exceptions=True,
            )
        self.state = ServerState.STOPPING

    @staticmethod
    async def _readChunks(input_stream: BinaryIO):
        try:
            input_stream.fileno()
        except (AttributeError, io.UnsupportedOperation):
            while True:
                chunk = input_stream.read(65_536)
                if not chunk:
                    return
                yield chunk
            return
        loop = asyncio.get_running_loop()
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        transport, _ = await loop.connect_read_pipe(
            lambda: protocol,
            input_stream,
        )
        try:
            while True:
                chunk = await reader.read(65_536)
                if not chunk:
                    return
                yield chunk
        finally:
            transport.close()

    async def processLine(self, raw_line: bytes) -> None:
        """Validate and schedule one newline-delimited JSON-RPC message."""
        if len(raw_line) > MAX_LINE_BYTES + 1:
            await self._writeError(None, -32700, "Parse error")
            return
        line = raw_line
        if line.endswith(b"\n"):
            line = line[:-1]
            if line.endswith(b"\r"):
                line = line[:-1]
        if not line:
            await self._writeError(None, -32700, "Parse error")
            return
        try:
            text = line.decode("utf-8", errors="strict")
            message = json.loads(text)
        except (UnicodeDecodeError, json.JSONDecodeError):
            await self._writeError(None, -32700, "Parse error")
            return
        if not isinstance(message, dict):
            await self._writeError(None, -32600, "Invalid Request")
            return
        is_notification = "id" not in message
        try:
            method, params = self._validateMessage(message)
        except JsonRpcError as error:
            if not is_notification:
                await self._writeError(
                    self._safeId(message.get("id")),
                    error.code,
                    error.message,
                    error.data,
                )
            return
        if is_notification:
            self._notification(method, params)
            return
        request_id = message.get("id")
        if not self._validId(request_id):
            await self._writeError(None, -32600, "Invalid Request")
            return
        task = asyncio.create_task(
            self._handleRequest(request_id, method, params),
            name=f"mcp-request-{request_id}",
        )
        self._tasks[request_id] = task
        task.add_done_callback(
            lambda _task, key=request_id: self._finishTask(key)
        )

    @staticmethod
    def _validId(value: object) -> bool:
        return (
            value is None
            or isinstance(value, str)
            or (isinstance(value, int) and not isinstance(value, bool))
        )

    @classmethod
    def _safeId(cls, value: object) -> object:
        return value if cls._validId(value) else None

    @staticmethod
    def _validateMessage(
        message: dict[str, object],
    ) -> tuple[str, dict[str, object]]:
        if message.get("jsonrpc") != "2.0":
            raise JsonRpcError(-32600, "Invalid Request")
        method = message.get("method")
        if not isinstance(method, str):
            raise JsonRpcError(-32600, "Invalid Request")
        params = message.get("params", {})
        if not isinstance(params, dict):
            raise JsonRpcError(-32602, "Invalid params")
        return method, params

    def _notification(
        self,
        method: str,
        params: dict[str, object],
    ) -> None:
        if method == "notifications/initialized":
            if self.state is ServerState.INITIALIZED:
                self.state = ServerState.RUNNING
            return
        if method == "notifications/cancelled":
            request_id = params.get("requestId")
            event = self._cancelled.get(request_id)
            if event is not None:
                event.set()
            return
        self.logger.debug("Ignoring unknown notification %s", method)

    def _finishTask(self, request_id: object) -> None:
        self._tasks.pop(request_id, None)
        self._cancelled.pop(request_id, None)

    async def _handleRequest(
        self,
        request_id: object,
        method: str,
        params: dict[str, object],
    ) -> None:
        try:
            result = await self._dispatch(request_id, method, params)
        except JsonRpcError as error:
            await self._writeError(
                request_id,
                error.code,
                error.message,
                error.data,
            )
        except Exception:
            self.logger.exception(
                "Unexpected MCP request failure for method %s",
                method,
            )
            await self._writeError(
                request_id,
                -32603,
                "Internal error",
            )
        else:
            await self._write({
                "jsonrpc": "2.0",
                "id": request_id,
                "result": result,
            })

    async def _dispatch(
        self,
        request_id: object,
        method: str,
        params: dict[str, object],
    ) -> object:
        if method == "initialize":
            return self._initialize(params)
        if method not in {"ping", "tools/list", "tools/call"}:
            raise JsonRpcError(-32601, "Method not found")
        if self.state not in {ServerState.INITIALIZED, ServerState.RUNNING}:
            raise JsonRpcError(-32002, "Server not initialized")
        if method == "ping":
            if params:
                raise JsonRpcError(-32602, "Invalid params")
            return {}
        if method == "tools/list":
            if set(params) - {"cursor", "_meta"}:
                raise JsonRpcError(-32602, "Invalid params")
            cursor = params.get("cursor")
            if cursor is not None and not isinstance(cursor, str):
                raise JsonRpcError(-32602, "Invalid params")
            meta = params.get("_meta")
            if meta is not None and not isinstance(meta, dict):
                raise JsonRpcError(-32602, "Invalid params")
            return {"tools": TOOLS}
        return await self._callTool(request_id, params)

    def _initialize(self, params: dict[str, object]) -> dict[str, object]:
        if self.state is not ServerState.STARTING:
            raise JsonRpcError(-32600, "Server is already initialized")
        requested = params.get("protocolVersion")
        if requested is not None and not isinstance(requested, str):
            raise JsonRpcError(-32602, "Invalid params")
        version = (
            requested
            if requested in SUPPORTED_PROTOCOL_VERSIONS
            else SUPPORTED_PROTOCOL_VERSIONS[0]
        )
        self.state = ServerState.INITIALIZED
        return {
            "protocolVersion": version,
            "capabilities": {"tools": {"listChanged": False}},
            "serverInfo": {
                "name": "firefox-rdp-mcp",
                "version": "0.1.0",
            },
        }

    async def _callTool(
        self,
        request_id: object,
        params: dict[str, object],
    ) -> dict[str, object]:
        allowed = {"name", "arguments", "_meta"}
        if set(params) - allowed:
            raise JsonRpcError(-32602, "Invalid params")
        meta = params.get("_meta")
        if meta is not None and not isinstance(meta, dict):
            raise JsonRpcError(-32602, "Invalid params")
        name = params.get("name")
        if not isinstance(name, str):
            raise JsonRpcError(-32602, "Invalid params")
        arguments = params.get("arguments", {})
        event = asyncio.Event()
        self._cancelled[request_id] = event
        try:
            return await self.tools.call(name, arguments, event.is_set)
        except InvalidArgumentsError as error:
            raise JsonRpcError(-32602, str(error)) from error
        except CancelledError as error:
            return errorResult(error)

    async def _writeError(
        self,
        request_id: object,
        code: int,
        message: str,
        data: object | None = None,
    ) -> None:
        error: dict[str, object] = {"code": code, "message": message}
        if data is not None:
            error["data"] = data
        await self._write({
            "jsonrpc": "2.0",
            "id": request_id,
            "error": error,
        })

    async def _write(self, message: dict[str, object]) -> None:
        payload = (
            json.dumps(
                message,
                ensure_ascii=False,
                separators=(",", ":"),
            ).encode("utf-8")
            + b"\n"
        )
        async with self._output_lock:
            self.output.write(payload)
            self.output.flush()


async def runServer(
    tools: ToolHandlers,
    input_stream: BinaryIO,
    output_stream: BinaryIO,
) -> None:
    """Construct and run an MCP server on supplied binary streams."""
    server = McpServer(tools, output_stream)
    await server.run(input_stream)