"""
MCP Transport Layer - Communication with MCP servers

Supports:
- STDIO transport (subprocess communication)
- SSE transport (Server-Sent Events)
- HTTP transport (REST-based)
"""

import asyncio
import json
import os
from abc import ABC, abstractmethod
from typing import Optional

from groknroll.mcp.types import MCPMessage


class MCPTransport(ABC):
    """Abstract base class for MCP transport implementations"""

    @abstractmethod
    async def connect(self) -> None:
        """Establish connection to the MCP server"""
        pass

    @abstractmethod
    async def disconnect(self) -> None:
        """Close connection to the MCP server"""
        pass

    @abstractmethod
    async def send(self, message: MCPMessage) -> None:
        """Send a message to the MCP server"""
        pass

    @abstractmethod
    async def receive(self) -> MCPMessage:
        """Receive a message from the MCP server"""
        pass

    @abstractmethod
    def is_connected(self) -> bool:
        """Check if transport is connected"""
        pass


class StdioTransport(MCPTransport):
    """
    STDIO-based MCP transport

    Communicates with MCP server via stdin/stdout using newline-delimited JSON.
    """

    def __init__(
        self,
        command: str,
        args: list[str],
        env: Optional[dict[str, str]] = None,
        cwd: Optional[str] = None,
    ):
        """
        Initialize STDIO transport

        Args:
            command: Command to run (e.g., "npx", "python", "uvx")
            args: Arguments to pass to the command
            env: Additional environment variables
            cwd: Working directory for the subprocess
        """
        self.command = command
        self.args = args
        self.env = env or {}
        self.cwd = cwd
        self._process: Optional[asyncio.subprocess.Process] = None
        self._read_lock = asyncio.Lock()
        self._write_lock = asyncio.Lock()

    async def connect(self) -> None:
        """Start the MCP server subprocess"""
        # Build environment
        process_env = os.environ.copy()
        process_env.update(self.env)

        # Start subprocess
        self._process = await asyncio.create_subprocess_exec(
            self.command,
            *self.args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=process_env,
            cwd=self.cwd,
        )

    async def disconnect(self) -> None:
        """Terminate the MCP server subprocess"""
        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()
            finally:
                self._process = None

    async def send(self, message: MCPMessage) -> None:
        """Send a message to the MCP server via stdin"""
        if not self._process or not self._process.stdin:
            raise RuntimeError("Transport not connected")

        async with self._write_lock:
            data = json.dumps(message.to_dict()) + "\n"
            self._process.stdin.write(data.encode())
            await self._process.stdin.drain()

    async def receive(self) -> MCPMessage:
        """Receive a message from the MCP server via stdout"""
        if not self._process or not self._process.stdout:
            raise RuntimeError("Transport not connected")

        async with self._read_lock:
            line = await self._process.stdout.readline()
            if not line:
                raise EOFError("MCP server closed connection")

            data = json.loads(line.decode().strip())
            return MCPMessage.from_dict(data)

    def is_connected(self) -> bool:
        """Check if subprocess is running"""
        return self._process is not None and self._process.returncode is None


class SSETransport(MCPTransport):
    """
    Server-Sent Events (SSE) based MCP transport

    Communicates with MCP server using HTTP SSE for receiving
    and HTTP POST for sending.
    """

    def __init__(
        self,
        url: str,
        headers: Optional[dict[str, str]] = None,
    ):
        """
        Initialize SSE transport

        Args:
            url: Base URL of the MCP server (e.g., "http://localhost:3000")
            headers: Additional HTTP headers
        """
        self.url = url.rstrip("/")
        self.headers = headers or {}
        self._connected = False
        self._message_queue: asyncio.Queue[MCPMessage] = asyncio.Queue()
        self._sse_task: Optional[asyncio.Task] = None

    async def connect(self) -> None:
        """Connect to SSE endpoint and start listening"""
        try:
            import aiohttp
        except ImportError:
            raise ImportError("aiohttp required for SSE transport: pip install aiohttp")

        self._session = aiohttp.ClientSession(headers=self.headers)

        # Start SSE listener
        self._sse_task = asyncio.create_task(self._sse_listener())
        self._connected = True

    async def _sse_listener(self) -> None:
        """Listen for SSE events"""
        try:
            import aiohttp
        except ImportError:
            return

        sse_url = f"{self.url}/sse"

        async with self._session.get(sse_url) as response:
            async for line in response.content:
                line = line.decode().strip()
                if line.startswith("data: "):
                    data = json.loads(line[6:])
                    message = MCPMessage.from_dict(data)
                    await self._message_queue.put(message)

    async def disconnect(self) -> None:
        """Disconnect from SSE endpoint"""
        self._connected = False

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

        if hasattr(self, "_session"):
            await self._session.close()

    async def send(self, message: MCPMessage) -> None:
        """Send a message via HTTP POST"""
        if not self._connected:
            raise RuntimeError("Transport not connected")

        post_url = f"{self.url}/message"
        async with self._session.post(post_url, json=message.to_dict()) as response:
            if response.status != 200:
                raise RuntimeError(f"Failed to send message: {response.status}")

    async def receive(self) -> MCPMessage:
        """Receive a message from the SSE queue"""
        if not self._connected:
            raise RuntimeError("Transport not connected")

        return await self._message_queue.get()

    def is_connected(self) -> bool:
        """Check if connected to SSE endpoint"""
        return self._connected


class HTTPTransport(MCPTransport):
    """
    HTTP-based MCP transport

    Simple request/response pattern using HTTP POST.
    """

    def __init__(
        self,
        url: str,
        headers: Optional[dict[str, str]] = None,
    ):
        """
        Initialize HTTP transport

        Args:
            url: Base URL of the MCP server
            headers: Additional HTTP headers
        """
        self.url = url.rstrip("/")
        self.headers = headers or {}
        self._connected = False

    async def connect(self) -> None:
        """Initialize HTTP session"""
        try:
            import aiohttp
        except ImportError:
            raise ImportError("aiohttp required for HTTP transport: pip install aiohttp")

        self._session = aiohttp.ClientSession(headers=self.headers)
        self._connected = True

    async def disconnect(self) -> None:
        """Close HTTP session"""
        self._connected = False
        if hasattr(self, "_session"):
            await self._session.close()

    async def send(self, message: MCPMessage) -> None:
        """Send message via HTTP POST (response handled by receive)"""
        if not self._connected:
            raise RuntimeError("Transport not connected")

        # Store the message for paired receive
        self._pending_message = message

    async def receive(self) -> MCPMessage:
        """Send pending message and receive response"""
        if not self._connected:
            raise RuntimeError("Transport not connected")

        if not hasattr(self, "_pending_message"):
            raise RuntimeError("No pending message to send")

        message = self._pending_message
        del self._pending_message

        async with self._session.post(self.url, json=message.to_dict()) as response:
            if response.status != 200:
                raise RuntimeError(f"HTTP request failed: {response.status}")

            data = await response.json()
            return MCPMessage.from_dict(data)

    def is_connected(self) -> bool:
        """Check if HTTP session is active"""
        return self._connected


def create_transport(
    transport_type: str,
    command: Optional[str] = None,
    args: Optional[list[str]] = None,
    env: Optional[dict[str, str]] = None,
    url: Optional[str] = None,
    headers: Optional[dict[str, str]] = None,
    cwd: Optional[str] = None,
) -> MCPTransport:
    """
    Create an MCP transport based on type

    Args:
        transport_type: "stdio", "sse", or "http"
        command: Command for stdio transport
        args: Arguments for stdio transport
        env: Environment variables for stdio transport
        url: URL for SSE/HTTP transport
        headers: HTTP headers for SSE/HTTP transport
        cwd: Working directory for stdio transport

    Returns:
        MCPTransport instance
    """
    if transport_type == "stdio":
        if not command:
            raise ValueError("command required for stdio transport")
        return StdioTransport(command, args or [], env, cwd)
    elif transport_type == "sse":
        if not url:
            raise ValueError("url required for SSE transport")
        return SSETransport(url, headers)
    elif transport_type == "http":
        if not url:
            raise ValueError("url required for HTTP transport")
        return HTTPTransport(url, headers)
    else:
        raise ValueError(f"Unknown transport type: {transport_type}")
