"""A2A (Agent-to-Agent) protocol adapter for Skytale encrypted channels.

Maps Google's A2A protocol to Skytale's MLS-encrypted channel infrastructure
using JSON-RPC 2.0 message framing.  Each A2A ``contextId`` becomes a Skytale
channel at ``org/a2a/{context_id}``, and JSON-RPC 2.0 messages are wrapped in
an :class:`~skytale_sdk.envelope.Envelope` tagged with ``Protocol.A2A`` and
transmitted through MLS-encrypted channels.

Two main classes are provided:

- :class:`A2AEncryptedTransport` — async JSON-RPC 2.0 transport (primary API).
  Provides ``read()``/``write()``/``close()`` for sending and receiving
  JSON-RPC 2.0 messages over encrypted A2A channels.

- :class:`SkytaleA2AAdapter` — synchronous task-based adapter (convenience API).
  Higher-level helper that manages context lifecycle and message send/receive
  with automatic JSON-RPC 2.0 envelope wrapping.

The adapter works with plain Python dicts today; an ``a2a-sdk`` dependency is
**not** required.  If/when we add direct ``a2a-sdk`` type support, it will be
imported lazily inside the methods that need it.

Usage (transport)::

    import asyncio
    from skytale_sdk.channels import SkytaleChannelManager
    from skytale_sdk.integrations._a2a import A2AEncryptedTransport

    mgr = SkytaleChannelManager(identity=b"agent-1")
    mgr.create("org/a2a/ctx-1")
    transport = A2AEncryptedTransport(mgr, "org/a2a/ctx-1")

    async def main():
        await transport.write({
            "jsonrpc": "2.0",
            "method": "message/send",
            "id": 1,
            "params": {"contextId": "ctx-1", "parts": [{"type": "text", "text": "hi"}]},
        })
        response = await transport.read()
        print(response)
        await transport.close()

    asyncio.run(main())

Usage (adapter)::

    from skytale_sdk.channels import SkytaleChannelManager
    from skytale_sdk.integrations._a2a import SkytaleA2AAdapter

    mgr = SkytaleChannelManager(identity=b"my-agent")
    adapter = SkytaleA2AAdapter(mgr, agent_id="agent-007")

    adapter.create_context("ctx-1")
    adapter.send_message("ctx-1", [{"type": "text", "text": "Hello!"}])
    msgs = adapter.receive_messages("ctx-1")
"""

from __future__ import annotations

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

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

logger = logging.getLogger(__name__)


# ======================================================================
# JSON-RPC 2.0 helpers
# ======================================================================

def _jsonrpc_request(
    method: str,
    params: Optional[Dict[str, Any]] = None,
    request_id: Optional[Any] = None,
) -> Dict[str, Any]:
    """Build a JSON-RPC 2.0 request dict.

    If *request_id* is ``None`` a random UUID is generated so every
    request is uniquely identifiable.

    Args:
        method: The JSON-RPC method name (e.g. ``"message/send"``).
        params: Optional parameter dict.
        request_id: Caller-specified request id.  Defaults to a UUID4 string.

    Returns:
        A JSON-RPC 2.0 request dict.

    Example::

        >>> req = _jsonrpc_request("message/send", {"text": "hi"})
        >>> req["jsonrpc"]
        '2.0'
    """
    msg: Dict[str, Any] = {
        "jsonrpc": "2.0",
        "method": method,
        "id": request_id if request_id is not None else str(uuid.uuid4()),
    }
    if params is not None:
        msg["params"] = params
    return msg


def _jsonrpc_response(
    request_id: Any,
    result: Any,
) -> Dict[str, Any]:
    """Build a JSON-RPC 2.0 success response dict.

    Args:
        request_id: The ``id`` from the originating request.
        result: The result payload.

    Returns:
        A JSON-RPC 2.0 response dict.

    Example::

        >>> resp = _jsonrpc_response(1, {"status": "ok"})
        >>> "error" not in resp
        True
    """
    return {"jsonrpc": "2.0", "id": request_id, "result": result}


def _jsonrpc_error(
    request_id: Any,
    code: int,
    message: str,
    data: Optional[Any] = None,
) -> Dict[str, Any]:
    """Build a JSON-RPC 2.0 error response dict.

    Args:
        request_id: The ``id`` from the originating request (may be ``None``
            for parse errors).
        code: Integer error code.
        message: Human-readable error description.
        data: Optional additional error data.

    Returns:
        A JSON-RPC 2.0 error response dict.

    Example::

        >>> err = _jsonrpc_error(1, -32600, "Invalid Request")
        >>> err["error"]["code"]
        -32600
    """
    err: Dict[str, Any] = {"code": code, "message": message}
    if data is not None:
        err["data"] = data
    return {"jsonrpc": "2.0", "id": request_id, "error": err}


def _jsonrpc_notification(
    method: str,
    params: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Build a JSON-RPC 2.0 notification dict (no ``id`` field).

    Notifications are fire-and-forget messages that do not expect a response.

    Args:
        method: The JSON-RPC method name.
        params: Optional parameter dict.

    Returns:
        A JSON-RPC 2.0 notification dict (no ``id`` key).

    Example::

        >>> n = _jsonrpc_notification("task/progress", {"percent": 50})
        >>> "id" not in n
        True
    """
    msg: Dict[str, Any] = {"jsonrpc": "2.0", "method": method}
    if params is not None:
        msg["params"] = params
    return msg


# ======================================================================
# A2AEncryptedTransport — async JSON-RPC 2.0 transport
# ======================================================================

class A2AEncryptedTransport:
    """A2A JSON-RPC 2.0 transport over MLS-encrypted Skytale channels.

    Replaces plaintext HTTP with end-to-end encrypted messaging for A2A
    protocol communication.  Each JSON-RPC 2.0 message is wrapped in an
    :class:`Envelope` with :attr:`Protocol.A2A`.

    This transport follows the same interface pattern as
    :class:`~skytale_sdk.integrations._mcp_transport.SkytaleTransport` so it
    can be used as a drop-in encrypted channel for A2A clients and servers.

    Args:
        manager: A :class:`SkytaleChannelManager` with the channel already
            created or joined.
        channel: Channel name (typically ``org/a2a/{context_id}``).
        agent_id: Optional agent identifier embedded in outgoing metadata.

    Example::

        mgr = SkytaleChannelManager(identity=b"agent-1")
        mgr.create("org/a2a/ctx-1")
        transport = A2AEncryptedTransport(mgr, "org/a2a/ctx-1")

        # In an async context:
        await transport.write({
            "jsonrpc": "2.0",
            "method": "message/send",
            "id": 1,
            "params": {"parts": [{"type": "text", "text": "Hello"}]},
        })
        msg = await transport.read()
        await transport.close()
    """

    def __init__(
        self,
        manager: SkytaleChannelManager,
        channel: str,
        agent_id: Optional[str] = None,
    ) -> None:
        self._manager = manager
        self._channel = channel
        self._agent_id = agent_id
        self._closed = False
        # Buffer for A2A 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] = []

    @property
    def channel(self) -> str:
        """The Skytale channel this transport is bound to.

        Example::

            >>> transport.channel
            'org/a2a/ctx-1'
        """
        return self._channel

    @property
    def closed(self) -> bool:
        """Whether the transport has been closed.

        Example::

            >>> transport.closed
            False
        """
        return self._closed

    async def read(self) -> Dict[str, Any]:
        """Read the next A2A JSON-RPC 2.0 message from the encrypted channel.

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

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

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            msg = await transport.read()
            if "method" in msg:
                print(f"request: {msg['method']}")
            elif "result" in msg:
                print(f"result for id={msg['id']}")
        """
        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_running_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 A2A-protocol envelopes only; other protocols sharing
            # the same channel are silently ignored.
            a2a_messages: List[Dict] = []
            for env in envelopes:
                if env.protocol is not Protocol.A2A:
                    logger.debug(
                        "ignoring non-A2A envelope (protocol=%s) on channel %s",
                        env.protocol.value,
                        self._channel,
                    )
                    continue
                try:
                    a2a_messages.append(json.loads(env.payload))
                except json.JSONDecodeError:
                    logger.warning(
                        "dropping A2A envelope with invalid JSON payload "
                        "(%d bytes) on channel %s",
                        len(env.payload),
                        self._channel,
                    )

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

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

    async def write(self, message: Dict[str, Any]) -> None:
        """Write an A2A JSON-RPC 2.0 message to the encrypted channel.

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

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

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            await transport.write({
                "jsonrpc": "2.0",
                "method": "message/send",
                "id": 1,
                "params": {"parts": [{"type": "text", "text": "Hello"}]},
            })
        """
        if self._closed:
            raise RuntimeError("transport closed")

        payload = json.dumps(message, separators=(",", ":")).encode("utf-8")

        metadata: Optional[Dict[str, str]] = None
        if self._agent_id:
            metadata = {"agent_id": self._agent_id}

        envelope = Envelope(
            protocol=Protocol.A2A,
            content_type="application/json",
            payload=payload,
            metadata=metadata,
        )

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

        logger.debug(
            "wrote A2A JSON-RPC message (%d bytes payload) to channel %s",
            len(payload),
            self._channel,
        )

    async def send_request(
        self,
        method: str,
        params: Optional[Dict[str, Any]] = None,
        request_id: Optional[Any] = None,
    ) -> str:
        """Build and send a JSON-RPC 2.0 request.

        Convenience wrapper around :meth:`write` that constructs a properly
        formed JSON-RPC 2.0 request and returns the assigned request id.

        Args:
            method: The A2A method name (e.g. ``"message/send"``,
                ``"tasks/get"``).
            params: Optional parameter dict.
            request_id: Caller-specified request id.  Defaults to a UUID4.

        Returns:
            The request id (string) so the caller can correlate responses.

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            rid = await transport.send_request("message/send", {
                "parts": [{"type": "text", "text": "Hello"}],
            })
        """
        req = _jsonrpc_request(method, params, request_id)
        await self.write(req)
        return req["id"]

    async def send_notification(
        self,
        method: str,
        params: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Build and send a JSON-RPC 2.0 notification (no response expected).

        Args:
            method: The A2A method name.
            params: Optional parameter dict.

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            await transport.send_notification("tasks/cancel", {"taskId": "t1"})
        """
        await self.write(_jsonrpc_notification(method, params))

    async def send_result(
        self,
        request_id: Any,
        result: Any,
    ) -> None:
        """Build and send a JSON-RPC 2.0 success response.

        Args:
            request_id: The ``id`` from the request being answered.
            result: The result payload.

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            await transport.send_result(1, {"status": "completed"})
        """
        await self.write(_jsonrpc_response(request_id, result))

    async def send_error(
        self,
        request_id: Any,
        code: int,
        message: str,
        data: Optional[Any] = None,
    ) -> None:
        """Build and send a JSON-RPC 2.0 error response.

        Args:
            request_id: The ``id`` from the request being answered.
            code: Integer error code (e.g. ``-32600`` for invalid request).
            message: Human-readable error description.
            data: Optional additional error data.

        Raises:
            RuntimeError: If the transport has been closed.

        Example::

            await transport.send_error(1, -32601, "Method not found")
        """
        await self.write(_jsonrpc_error(request_id, code, message, data))

    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("A2A transport closed for channel %s", self._channel)


# ======================================================================
# SkytaleA2AAdapter — synchronous convenience wrapper
# ======================================================================

class SkytaleA2AAdapter:
    """Maps A2A protocol concepts to Skytale encrypted channels.

    Each A2A context maps to a Skytale channel at ``org/a2a/{context_id}``.
    Messages are serialized as JSON-RPC 2.0 ``message/send`` requests,
    wrapped in an :class:`Envelope` with ``Protocol.A2A``, and sent through
    MLS-encrypted channels.

    Args:
        manager: A :class:`SkytaleChannelManager` that owns the underlying
            encrypted channels.
        agent_id: Identifier for this agent, embedded in every outgoing A2A
            message so receivers know who sent it.

    Example::

        from skytale_sdk.channels import SkytaleChannelManager
        from skytale_sdk.integrations._a2a import SkytaleA2AAdapter

        mgr = SkytaleChannelManager(identity=b"a2a-agent")
        adapter = SkytaleA2AAdapter(mgr, agent_id="agent-42")
        adapter.create_context("research-session")
        adapter.send_message("research-session", [
            {"type": "text", "text": "Analysis complete."},
        ])
    """

    def __init__(self, manager: SkytaleChannelManager, agent_id: str) -> None:
        self._manager = manager
        self._agent_id = agent_id

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _channel_name(self, context_id: str) -> str:
        """Convert an A2A context ID to a Skytale channel name.

        A2A contexts live under the ``org/a2a/`` namespace so they do not
        collide with channels created by other protocols or directly by
        users.

        Args:
            context_id: The A2A context identifier.

        Returns:
            A Skytale channel name in ``org/a2a/{context_id}`` format.

        Example::

            >>> adapter._channel_name("ctx-1")
            'org/a2a/ctx-1'
        """
        return f"org/a2a/{context_id}"

    # ------------------------------------------------------------------
    # Context lifecycle
    # ------------------------------------------------------------------

    def create_context(self, context_id: str) -> None:
        """Create a new A2A context backed by a Skytale encrypted channel.

        The underlying MLS group is created immediately and a background
        listener thread starts buffering incoming messages.

        Args:
            context_id: Unique identifier for the A2A conversation context.

        Raises:
            RuntimeError: If channel creation fails (e.g. invalid name or
                relay unreachable).

        Example::

            adapter.create_context("research-session")
        """
        channel = self._channel_name(context_id)
        self._manager.create(channel)
        logger.info(
            "created A2A context %s on channel %s",
            context_id,
            channel,
        )

    def join_context(self, context_id: str, welcome: bytes) -> None:
        """Join an existing A2A context using an MLS Welcome message.

        Another agent must have already created the context and generated
        the Welcome via :meth:`SkytaleChannelManager.add_member`.

        Args:
            context_id: The A2A context identifier to join.
            welcome: MLS Welcome message bytes from the context creator.

        Raises:
            RuntimeError: If the Welcome is invalid or joining fails.

        Example::

            # On the joining agent, after receiving *welcome* out-of-band:
            adapter.join_context("research-session", welcome_bytes)
        """
        channel = self._channel_name(context_id)
        self._manager.join(channel, welcome)
        logger.info(
            "joined A2A context %s on channel %s",
            context_id,
            channel,
        )

    # ------------------------------------------------------------------
    # Messaging
    # ------------------------------------------------------------------

    def send_message(self, context_id: str, parts: List[Dict]) -> None:
        """Send an A2A message with the given parts.

        Builds a JSON-RPC 2.0 ``message/send`` request containing the A2A
        message parts, wraps it in a :class:`~skytale_sdk.envelope.Envelope`
        with ``Protocol.A2A``, and sends it over the MLS-encrypted channel.

        The request is assigned a random ``id`` (UUID4).  The ``agentId``
        and ``contextId`` fields are populated automatically inside the
        params object.

        Args:
            context_id: A2A context to send the message on.
            parts: List of A2A message parts.  Each part is a dict, e.g.
                ``{"type": "text", "text": "Hello"}``.

        Raises:
            KeyError: If the context has not been created or joined.

        Example::

            adapter.send_message("ctx-1", [
                {"type": "text", "text": "Hello from the research agent."},
            ])
        """
        request_id = str(uuid.uuid4())
        jsonrpc_message = _jsonrpc_request(
            method="message/send",
            params={
                "message": {
                    "messageId": str(uuid.uuid4()),
                    "contextId": context_id,
                    "agentId": self._agent_id,
                    "role": "agent",
                    "parts": parts,
                },
            },
            request_id=request_id,
        )

        payload = json.dumps(jsonrpc_message, separators=(",", ":")).encode("utf-8")

        envelope = Envelope(
            protocol=Protocol.A2A,
            content_type="application/json",
            payload=payload,
            metadata={"context_id": context_id},
        )

        channel = self._channel_name(context_id)
        self._manager.send_envelope(channel, envelope)
        logger.debug(
            "sent A2A message (request %s) on context %s (%d bytes)",
            request_id,
            context_id,
            len(payload),
        )

    def receive_messages(
        self, context_id: str, timeout: float = 5.0
    ) -> List[Dict]:
        """Receive A2A messages from a context.

        Drains all buffered envelopes from the channel, filters for
        ``Protocol.A2A``, deserializes their JSON payloads, and returns
        them as plain dicts.  Both JSON-RPC 2.0 requests/responses and
        plain A2A message dicts are accepted.

        Non-A2A envelopes (e.g. raw strings or MCP messages on a bridged
        channel) are silently skipped.  Envelopes whose payloads are not
        valid JSON are logged as warnings and skipped.

        Args:
            context_id: A2A context to read from.
            timeout: Maximum seconds to wait for messages.  ``0`` returns
                immediately.

        Returns:
            List of A2A message dicts.  Empty if no A2A messages arrive
            before the timeout expires.

        Raises:
            KeyError: If the context has not been created or joined.
            RuntimeError: If the background listener thread has died.

        Example::

            msgs = adapter.receive_messages("ctx-1", timeout=2.0)
            for msg in msgs:
                for part in msg.get("parts", []):
                    if part.get("type") == "text":
                        print(part["text"])
        """
        channel = self._channel_name(context_id)
        envelopes = self._manager.receive_envelopes(channel, timeout=timeout)

        messages: List[Dict] = []
        for env in envelopes:
            if env.protocol is not Protocol.A2A:
                # Not an A2A message — skip without error so mixed-protocol
                # channels degrade gracefully.
                continue
            try:
                msg = json.loads(env.payload)
            except (json.JSONDecodeError, UnicodeDecodeError) as exc:
                logger.warning(
                    "skipping A2A envelope with invalid JSON payload: %s", exc
                )
                continue
            messages.append(msg)

        logger.debug(
            "received %d A2A message(s) on context %s",
            len(messages),
            context_id,
        )
        return messages

    # ------------------------------------------------------------------
    # Discovery
    # ------------------------------------------------------------------

    def agent_card_extension(self) -> Dict:
        """Return Skytale channel metadata for embedding in an A2A AgentCard.

        The returned dict can be placed under an ``extensions.skytale`` key
        in an A2A AgentCard so that other agents know how to reach this
        agent over a Skytale encrypted channel.

        Returns:
            Dict with keys ``"agent_id"``, ``"channel_prefix"``, and
            ``"active_contexts"`` (list of context IDs currently open).

        Example::

            card_ext = adapter.agent_card_extension()
            # card_ext == {
            #     "agent_id": "agent-42",
            #     "channel_prefix": "org/a2a/",
            #     "active_contexts": ["ctx-1", "ctx-2"],
            # }
        """
        # Derive active context IDs by inspecting channels that match
        # the A2A namespace prefix.
        prefix = "org/a2a/"
        active = [
            ch[len(prefix):]
            for ch in self._manager.list_channels()
            if ch.startswith(prefix)
        ]

        return {
            "agent_id": self._agent_id,
            "channel_prefix": prefix,
            "active_contexts": active,
        }
