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

Maps Google's A2A protocol concepts to Skytale's MLS-encrypted channel
infrastructure.  Each A2A ``contextId`` becomes a Skytale channel at
``org/a2a/{context_id}``, and A2A ``Message`` objects (with ``parts``) are
JSON-serialized, wrapped in an :class:`~skytale_sdk.envelope.Envelope` tagged
with ``Protocol.A2A``, and transmitted through MLS-encrypted channels.

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::

    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 json
import logging
import uuid
from typing import Dict, List, Optional

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

logger = logging.getLogger(__name__)


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, 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 object conforming to the A2A ``Message`` shape, wraps
        it in a :class:`~skytale_sdk.envelope.Envelope` with ``Protocol.A2A``,
        and sends it over the MLS-encrypted channel.

        The message is assigned a random ``messageId`` (UUID4).  The
        ``agentId`` and ``contextId`` fields are populated automatically.

        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."},
            ])
        """
        message_id = str(uuid.uuid4())
        # Build an A2A-shaped message dict.  We intentionally keep this as
        # a plain dict (rather than importing a2a-sdk types) so the adapter
        # has zero optional dependencies today.
        a2a_message: Dict = {
            "messageId": message_id,
            "contextId": context_id,
            "agentId": self._agent_id,
            "role": "agent",
            "parts": parts,
        }

        payload = json.dumps(a2a_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 %s on context %s (%d bytes)",
            message_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.

        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,
        }
