"""Streamable HTTP transport for MCP servers (preferred).

This transport uses a single HTTP endpoint that supports:
- POST for JSON-RPC requests/notifications/responses
- GET for optional SSE streams carrying notifications and/or responses

It also supports session management via `MCP-Session-Id` and protocol version
header forwarding.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any

import httpx

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

logger = logging.getLogger("mcpbridge")


class StreamableHttpTransport(BaseTransport):
    """Streamable HTTP MCP transport implementation."""

    def __init__(self, url: str, headers: dict[str, str] | None = None, timeout: int = 30) -> None:
        self._url = url.rstrip("/")
        self._headers = headers or {}
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

        self._connected = False
        self._sse_task: asyncio.Task[None] | None = None
        self._pending: dict[str | int, asyncio.Future[dict[str, Any]]] = {}

        self._session_id: str | None = None
        self._last_event_id: str | None = None

        self.notifications_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

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

    async def connect(self) -> None:
        if self._connected:
            return
        self._client = httpx.AsyncClient(timeout=self._timeout)
        self._connected = True
        self._sse_task = asyncio.create_task(self._background_sse_listener(), name="mcpbridge-streamable-http-sse")

    async def close(self) -> None:
        self._connected = False
        if self._sse_task:
            self._sse_task.cancel()
        for fut in list(self._pending.values()):
            if not fut.done():
                fut.set_exception(TransportDisconnectedError("Transport disconnected"))
        self._pending.clear()
        if self._client:
            await self._client.aclose()
        self._client = None

    async def send_notification(self, method: str, params: dict[str, Any]) -> None:
        payload: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            payload["params"] = params
        await self._post(payload)

    async def send(self, message: dict[str, Any]) -> dict[str, Any]:
        if not self._client:
            raise TransportDisconnectedError("Transport not connected")

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

        loop = asyncio.get_running_loop()
        fut: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[msg_id] = fut
        try:
            resp = await self._post(message)
            # If _post returns sync JSON response, it will have resolved fut.
            # If _post returns SSE stream, it will resolve fut itself before returning.
            if not fut.done():
                await asyncio.wait_for(fut, timeout=self._timeout)
            return fut.result()
        except asyncio.TimeoutError as e:
            self._pending.pop(msg_id, None)
            raise TimeoutError("Transport timed out") from e
        finally:
            self._pending.pop(msg_id, None)

    async def _post(self, payload: dict[str, Any]) -> None:
        assert self._client is not None
        headers: dict[str, str] = {
            "Accept": "application/json, text/event-stream",
            **self._headers,
        }
        if self._session_id:
            # Spec uses `Mcp-Session-Id` but header matching is case-insensitive.
            headers["MCP-Session-Id"] = self._session_id

        try:
            resp = await self._client.post(self._url, json=payload, headers=headers)
        except Exception as e:
            raise TransportConnectionError("Failed to POST to MCP server") from e

        # Capture session id when provided (init phase typically).
        session_hdr = resp.headers.get("MCP-Session-Id") or resp.headers.get("Mcp-Session-Id")
        if session_hdr:
            self._session_id = session_hdr

        ct = resp.headers.get("content-type", "")
        if resp.status_code in {200, 201}:
            if "application/json" in ct:
                data = resp.json()
                msg_id = data.get("id")
                if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
                    self._pending[msg_id].set_result(data)
                return
            if "text/event-stream" in ct:
                # Inline SSE: parse until we get a response.
                await self._parse_sse_stream(resp, expect_id=payload.get("id"))
                return
            # Sometimes servers return JSON without proper content-type
            try:
                data = resp.json()
                msg_id = data.get("id")
                if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
                    self._pending[msg_id].set_result(data)
                return
            except Exception:
                raise MCPProtocolError("Unexpected response content-type from MCP server")

        if resp.status_code == 202:
            # Response will arrive on SSE stream.
            return

        if resp.status_code == 404 and self._session_id:
            # Session expired: start new session by clearing the header and retry once.
            self._session_id = None
            headers.pop("MCP-Session-Id", None)
            resp2 = await self._client.post(self._url, json=payload, headers=headers)
            session_hdr = resp2.headers.get("MCP-Session-Id") or resp2.headers.get("Mcp-Session-Id")
            if session_hdr:
                self._session_id = session_hdr
            if resp2.status_code >= 400:
                raise TransportConnectionError(f"MCP server HTTP error {resp2.status_code}: {resp2.text}")
            if "application/json" in resp2.headers.get("content-type", ""):
                data = resp2.json()
                msg_id = data.get("id")
                if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
                    self._pending[msg_id].set_result(data)
                return
            if "text/event-stream" in resp2.headers.get("content-type", ""):
                # Parse response SSE from retry response.
                await self._parse_sse_stream(resp2, expect_id=payload.get("id"))
                return

        raise TransportConnectionError(f"MCP server HTTP error: {resp.status_code}: {resp.text}")

    async def _parse_sse_stream(self, response: httpx.Response, expect_id: str | int | None) -> None:
        """Parse an SSE response body until the expected JSON-RPC response arrives."""
        # httpx.Response is already fully received for non-streaming POST.
        # When content-type is text/event-stream, we parse the body line-by-line.
        buffer: list[str] = []
        last_event_id: str | None = None
        async for line in response.aiter_lines():
            if line is None:
                continue
            line = line.strip("\r")
            if line.startswith("id:"):
                last_event_id = line[len("id:") :].strip()
            if line == "":
                if not buffer:
                    continue
                data_str = "\n".join(buffer)
                buffer = []
                if not data_str.strip():
                    continue
                try:
                    msg = json.loads(data_str)
                except Exception as e:
                    raise MCPProtocolError("Malformed JSON from MCP SSE") from e
                if expect_id is not None and msg.get("id") == expect_id:
                    msg_id = msg.get("id")
                    if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
                        self._pending[msg_id].set_result(msg)
                    return
                await self._dispatch_incoming(msg)
                # If no expect_id, just dispatch; if we got response with id, dispatch and continue.
                continue

            if line.startswith("data:"):
                buffer.append(line[len("data:") :].lstrip())

        if expect_id is not None and expect_id in self._pending and not self._pending[expect_id].done():
            raise TransportDisconnectedError("SSE ended before response arrived")

    async def _background_sse_listener(self) -> None:
        """Listen for server notifications (and potentially async JSON-RPC messages)."""
        assert self._client is not None
        while self._connected:
            try:
                headers: dict[str, str] = {"Accept": "text/event-stream", **self._headers}
                if self._session_id:
                    headers["MCP-Session-Id"] = self._session_id
                if self._last_event_id:
                    headers["Last-Event-ID"] = self._last_event_id

                async with self._client.stream("GET", self._url, headers=headers) as resp:
                    if resp.status_code == 405:
                        # Server does not offer SSE at this endpoint.
                        return
                    if resp.status_code >= 400:
                        raise TransportConnectionError(f"SSE endpoint error: {resp.status_code}")

                    buffer: list[str] = []
                    retry_ms: int | None = None
                    last_event_id: str | None = None
                    async for line in resp.aiter_lines():
                        if line is None:
                            continue
                        line = line.strip("\r")
                        if line == "":
                            if not buffer:
                                continue
                            data_str = "\n".join(buffer)
                            buffer = []
                            if not data_str.strip():
                                continue
                            try:
                                msg = json.loads(data_str)
                            except Exception as e:
                                raise MCPProtocolError("Malformed JSON from MCP SSE") from e
                            await self._dispatch_incoming(msg)
                            self._last_event_id = last_event_id or self._last_event_id
                            continue

                        if line.startswith("id:"):
                            last_event_id = line[len("id:") :].strip() or None
                        if line.startswith("retry:"):
                            try:
                                retry_ms = int(line[len("retry:") :].strip())
                            except ValueError:
                                retry_ms = None
                        if line.startswith("data:"):
                            buffer.append(line[len("data:") :].lstrip())

                    # Stream ended; allow retry per SSE semantics.
                    if retry_ms:
                        await asyncio.sleep(retry_ms / 1000)
                    else:
                        await asyncio.sleep(0.5)
            except asyncio.CancelledError:
                return
            except Exception as e:
                if not self._connected:
                    return
                logger.warning("SSE listener disconnected (%s). Retrying...", e)
                await asyncio.sleep(0.5)

    async def _dispatch_incoming(self, msg: dict[str, Any]) -> None:
        msg_id = msg.get("id")
        if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
            self._pending[msg_id].set_result(msg)
            return
        method = msg.get("method")
        if msg_id is None and isinstance(method, str):
            await self.notifications_queue.put(msg)

