"""MCP JSON-RPC transport over MLS-encrypted Skytale channels.

Provides :class:`SkytaleTransport` — a drop-in transport for MCP clients and
servers that replaces plaintext HTTP/stdio with end-to-end encrypted messaging.
Each JSON-RPC message is serialized as JSON, wrapped in an :class:`Envelope`
tagged with :attr:`Protocol.MCP`, and sent through an MLS-encrypted channel.

This module is *not* the same as ``_mcp.py`` (which exposes Skytale operations
as MCP tools).  This module provides Skytale **as** the transport layer for
arbitrary MCP protocol messages.

Usage::

    import asyncio
    from skytale_sdk.channels import SkytaleChannelManager
    from skytale_sdk.integrations._mcp_transport import SkytaleTransport

    mgr = SkytaleChannelManager(identity=b"agent-1")
    mgr.create("org/ns/mcp-rpc")
    transport = SkytaleTransport(mgr, "org/ns/mcp-rpc")

    async def main():
        await transport.write({"jsonrpc": "2.0", "method": "ping", "id": 1})
        response = await transport.read()
        print(response)
        await transport.close()

    asyncio.run(main())
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Dict, List, Optional

from skytale_sdk.channels import SkytaleChannelManager
from skytale_sdk.envelope import Envelope, Protocol

logger = logging.getLogger(__name__)


class SkytaleTransport:
    """MCP JSON-RPC transport over an MLS-encrypted Skytale channel.

    Replaces plaintext HTTP/stdio with end-to-end encrypted messaging.
    Each JSON-RPC message is wrapped in an :class:`Envelope` with
    :attr:`Protocol.MCP`.

    Can be used as a drop-in transport for MCP clients and servers
    that support custom transports.

    Args:
        manager: A :class:`SkytaleChannelManager` with the channel already
            created or joined.
        channel: Channel name in ``org/namespace/service`` format.

    Example::

        mgr = SkytaleChannelManager(identity=b"agent-1")
        mgr.create("org/ns/mcp-rpc")
        transport = SkytaleTransport(mgr, "org/ns/mcp-rpc")

        # In an async context:
        await transport.write({"jsonrpc": "2.0", "method": "ping", "id": 1})
        msg = await transport.read()
        await transport.close()
    """

    def __init__(self, manager: SkytaleChannelManager, channel: str) -> None:
        self._manager = manager
        self._channel = channel
        self._closed = False
        # Buffer for MCP messages that arrived in a batch but haven't been
        # consumed yet — avoids dropping messages when receive_envelopes
        # returns multiple envelopes at once.
        self._pending: List[Dict] = []

    async def read(self) -> Dict:
        """Read the next MCP JSON-RPC message from the encrypted channel.

        Polls the underlying channel for :attr:`Protocol.MCP` envelopes,
        deserializes the first one's payload as JSON, and returns it.  If
        multiple MCP envelopes arrive in a single poll, the extras are queued
        internally and returned by subsequent ``read()`` calls.

        Returns:
            A ``dict`` representing the JSON-RPC message.

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            msg = await transport.read()
            print(msg["method"])  # e.g. "tools/list"
        """
        if self._closed:
            raise RuntimeError("transport closed")

        # Drain the internal queue first before hitting the channel.
        if self._pending:
            return self._pending.pop(0)

        loop = asyncio.get_event_loop()

        while True:
            if self._closed:
                raise RuntimeError("transport closed")

            # receive_envelopes is synchronous — run in executor to avoid
            # blocking the event loop.
            envelopes: List[Envelope] = await loop.run_in_executor(
                None,
                self._manager.receive_envelopes,
                self._channel,
                1.0,
            )

            # Filter to MCP-protocol envelopes only; other protocols sharing
            # the same channel are silently ignored.
            mcp_messages: List[Dict] = []
            for env in envelopes:
                if env.protocol is not Protocol.MCP:
                    logger.debug(
                        "ignoring non-MCP envelope (protocol=%s) on channel %s",
                        env.protocol.value,
                        self._channel,
                    )
                    continue
                try:
                    mcp_messages.append(json.loads(env.payload))
                except json.JSONDecodeError:
                    logger.warning(
                        "dropping MCP envelope with invalid JSON payload "
                        "(%d bytes) on channel %s",
                        len(env.payload),
                        self._channel,
                    )

            if mcp_messages:
                # Return the first message now; stash the rest for later reads.
                self._pending.extend(mcp_messages[1:])
                return mcp_messages[0]

            # No MCP messages this round — loop and poll again.

    async def write(self, message: Dict) -> None:
        """Write an MCP JSON-RPC message to the encrypted channel.

        Serializes *message* to JSON bytes, wraps it in an :class:`Envelope`
        tagged with :attr:`Protocol.MCP`, and sends it through the channel.

        Args:
            message: A JSON-serializable ``dict`` (typically a JSON-RPC
                request, response, or notification).

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            await transport.write({
                "jsonrpc": "2.0",
                "method": "tools/list",
                "id": 1,
            })
        """
        if self._closed:
            raise RuntimeError("transport closed")

        payload = json.dumps(message, separators=(",", ":")).encode("utf-8")
        envelope = Envelope(
            protocol=Protocol.MCP,
            content_type="application/json",
            payload=payload,
        )

        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            None,
            self._manager.send_envelope,
            self._channel,
            envelope,
        )

        logger.debug(
            "wrote MCP message (%d bytes payload) to channel %s",
            len(payload),
            self._channel,
        )

    async def close(self) -> None:
        """Close the transport.

        After closing, any subsequent :meth:`read` or :meth:`write` call will
        raise :class:`RuntimeError`.  Pending buffered messages are discarded.

        Example::

            await transport.close()
        """
        self._closed = True
        self._pending.clear()
        logger.debug("transport closed for channel %s", self._channel)
