import logging
import traceback
import uuid
from contextlib import asynccontextmanager
from typing import Literal

import anyio
import mcp.types as types
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.server.fastmcp import FastMCP as FastMCPBase
from mcp.shared.message import SessionMessage
from websockets.asyncio.server import ServerConnection, serve

from ..common.env import env

logger = logging.getLogger(__name__)


class BlaxelMcpServerTransport:
    """WebSocket server transport for MCP."""

    def __init__(self, port: int = 8080):
        """Initialize the WebSocket server transport.

        Args:
            port: The port to listen on (defaults to 8080 or BL_SERVER_PORT env var)
        """
        if env["BL_SERVER_PORT"] is not None:
            self.port = int(env["BL_SERVER_PORT"])
        else:
            self.port = port

        self.clients = {}
        self.server = None


    @asynccontextmanager
    async def websocket_server(self):
        """Create and run a WebSocket server for MCP communication."""
        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
        read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]

        write_stream: MemoryObjectSendStream[SessionMessage]
        write_stream_reader: MemoryObjectReceiveStream[SessionMessage]

        read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
        write_stream, write_stream_reader = anyio.create_memory_object_stream(0)

        async def handler(websocket: ServerConnection):
            client_id = str(uuid.uuid4())
            self.clients[client_id] = websocket
            logger.info(f"Client connected: {client_id}")

            try:
                async for message in websocket:
                    try:
                        msg = types.JSONRPCMessage.model_validate_json(message)
                        # Modify message ID to include client ID
                        if hasattr(msg, "id") and msg.id is not None:
                            original_id = msg.id
                            msg.id = f"{client_id}:{original_id}"
                        await read_stream_writer.send(SessionMessage(message=msg))
                    except Exception as exc:
                        logger.error(f"Failed to parse message: {exc}\n{traceback.format_exc()}")
                        await read_stream_writer.send(exc)
            except Exception as e:
                logger.error(f"WebSocket error: {e}")
            finally:
                if client_id in self.clients:
                    del self.clients[client_id]
                logger.info(f"Client disconnected: {client_id}")

        async def message_sender():
            async with write_stream_reader:
                async for session_message in write_stream_reader:
                    # Extract client ID from message ID
                    client_id = None
                    msg_id = None
                    message = session_message.message
                    if hasattr(message, "id") and message.id is not None:
                        parts = str(message.id).split(":", 1)
                        if len(parts) == 2:
                            client_id, msg_id = parts
                            # Restore original message ID
                            message.id = int(msg_id) if msg_id.isdigit() else msg_id

                    data = message.model_dump_json(by_alias=True, exclude_none=True)

                    if client_id and client_id in self.clients:
                        # Send to specific client with retry
                        websocket = self.clients[client_id]
                        success = await self._send_with_retry(websocket, data, client_id)
                        if not success:
                            # Remove client after failed retries
                            if client_id in self.clients:
                                del self.clients[client_id]
                                logger.info(f"Removed client {client_id} after failed message delivery")
                    else:
                        # Broadcast to all clients with retry
                        dead_clients = []
                        for cid, websocket in self.clients.items():
                            success = await self._send_with_retry(websocket, data, cid)
                            if not success:
                                dead_clients.append(cid)

                        # Clean up dead clients
                        for cid in dead_clients:
                            if cid in self.clients:
                                del self.clients[cid]
                                logger.info(f"Removed client {cid} after failed broadcast delivery")

        async with anyio.create_task_group() as tg:
            logger.info(f"Starting WebSocket Server on port {self.port}")
            async with serve(
                handler,
                "0.0.0.0",
                self.port,
                ping_timeout=None,   # Disable wait timeout
            ) as server:
                self.server = server
                tg.start_soon(message_sender)
                yield read_stream, write_stream


class FastMCP(FastMCPBase):
    def run(self, transport: Literal["stdio", "sse", "ws"] = "stdio") -> None:
        """Run the FastMCP server. Note this is a synchronous function.

        Args:
            transport: Transport protocol to use ("stdio" or "sse")
        """
        TRANSPORTS = Literal["stdio", "sse", "ws"]
        if transport not in TRANSPORTS.__args__:  # type: ignore
            raise ValueError(f"Unknown transport: {transport}")

        if transport == "stdio":
            anyio.run(self.run_stdio_async)
        elif transport == "ws":
            anyio.run(self.run_ws_async)
        else:
            anyio.run(self.run_sse_async)

    async def run_ws_async(self) -> None:
        async with BlaxelMcpServerTransport().websocket_server() as (read_stream, write_stream):
            await self._mcp_server.run(
                read_stream,
                write_stream,
                self._mcp_server.create_initialization_options(),
            )
