"""Cross-protocol bridge for translating between A2A, MCP, and SLIM.

SLIM is the universal transport -- following AGNTCY's pattern, N protocol
adapters compose freely without N*M bridges.  When translating between
non-SLIM protocols (e.g. A2A <-> MCP) we go direct for efficiency, but
every adapter can also round-trip through SLIM as the canonical
intermediary.

All messages stay MLS-encrypted in transit.  The bridge only sees
decrypted payloads *inside* the SDK process after the channel manager
has already decrypted them.

Usage::

    from skytale_sdk.channels import SkytaleChannelManager
    from skytale_sdk.bridge import ProtocolBridge
    from skytale_sdk.envelope import Protocol

    mgr = SkytaleChannelManager(identity=b"bridge-agent")
    mgr.create("org/a2a-agents")
    mgr.create("org/slim-agents")

    bridge = ProtocolBridge(mgr)
    bridge.bridge("org/a2a-agents", "org/slim-agents",
                  Protocol.A2A, Protocol.SLIM)
    # ... messages flow automatically ...
    bridge.stop()
"""

from __future__ import annotations

import base64
import json
import logging
import threading
import uuid
from typing import Callable, Dict, Optional, Tuple

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

logger = logging.getLogger(__name__)

# Translation dispatch key: (source_protocol, target_protocol)
_TranslationKey = Tuple[Protocol, Protocol]


class ProtocolBridge:
    """Translate messages between protocols through SLIM channels.

    Runs background threads that receive envelopes from a source channel,
    translate them to the target protocol, and send on the target channel.
    All messages stay MLS-encrypted -- the bridge translates between
    decrypted payloads only.

    Args:
        manager: Channel manager that owns the source and target channels.

    Example::

        mgr = SkytaleChannelManager(identity=b"bridge-agent")
        mgr.create("org/a2a-in")
        mgr.create("org/slim-out")

        bridge = ProtocolBridge(mgr)
        bridge.bridge("org/a2a-in", "org/slim-out",
                      Protocol.A2A, Protocol.SLIM)
        bridge.stop()
    """

    def __init__(self, manager: SkytaleChannelManager) -> None:
        self._manager = manager
        self._threads: list = []
        self._running = True

        # Dispatch table avoids a chain of if/elif in _translate.
        # Keys are (source, target) protocol pairs.
        self._translators: Dict[_TranslationKey, Callable[[Envelope], Envelope]] = {
            (Protocol.A2A, Protocol.SLIM): self._a2a_to_slim,
            (Protocol.SLIM, Protocol.A2A): self._slim_to_a2a,
            (Protocol.MCP, Protocol.SLIM): self._mcp_to_slim,
            (Protocol.SLIM, Protocol.MCP): self._slim_to_mcp,
            (Protocol.A2A, Protocol.MCP): self._a2a_to_mcp,
            (Protocol.MCP, Protocol.A2A): self._mcp_to_a2a,
            (Protocol.ANP, Protocol.SLIM): self._anp_to_slim,
            (Protocol.SLIM, Protocol.ANP): self._slim_to_anp,
            (Protocol.LMOS, Protocol.SLIM): self._lmos_to_slim,
            (Protocol.SLIM, Protocol.LMOS): self._slim_to_lmos,
            (Protocol.NLIP, Protocol.SLIM): self._nlip_to_slim,
            (Protocol.SLIM, Protocol.NLIP): self._slim_to_nlip,
        }

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def bridge(
        self,
        source_channel: str,
        target_channel: str,
        source_protocol: Protocol,
        target_protocol: Protocol,
    ) -> None:
        """Start bridging messages between two channels with translation.

        Spawns a daemon thread that continuously receives envelopes from
        *source_channel*, translates each to *target_protocol*, and sends
        the result on *target_channel*.

        Args:
            source_channel: Channel name to read from.
            target_channel: Channel name to write to.
            source_protocol: Protocol expected on the source channel.
            target_protocol: Protocol to translate into for the target channel.

        Example::

            bridge.bridge("org/mcp-tools", "org/slim-bus",
                          Protocol.MCP, Protocol.SLIM)
        """
        thread = threading.Thread(
            target=self._bridge_loop,
            args=(source_channel, target_channel, source_protocol, target_protocol),
            daemon=True,
            name=f"skytale-bridge-{source_channel}->{target_channel}",
        )
        self._threads.append(thread)
        thread.start()
        logger.info(
            "bridge started: %s (%s) -> %s (%s)",
            source_channel,
            source_protocol.value,
            target_channel,
            target_protocol.value,
        )

    def stop(self) -> None:
        """Signal all bridge threads to stop.

        Threads are daemons so they will not prevent interpreter exit.
        This method returns immediately after setting the stop flag.

        Example::

            bridge.stop()
        """
        self._running = False
        logger.info("bridge stop requested, %d threads signalled", len(self._threads))

    # ------------------------------------------------------------------
    # Internal: bridge loop
    # ------------------------------------------------------------------

    def _bridge_loop(
        self,
        source_channel: str,
        target_channel: str,
        source_protocol: Protocol,
        target_protocol: Protocol,
    ) -> None:
        """Poll source channel, translate, and forward.  Runs in a thread."""
        while self._running:
            try:
                envelopes = self._manager.receive_envelopes(
                    source_channel, timeout=1.0
                )
            except Exception:
                logger.exception(
                    "error receiving from %s, retrying", source_channel
                )
                continue

            for envelope in envelopes:
                try:
                    translated = self._translate(envelope, target_protocol)
                    self._manager.send_envelope(target_channel, translated)
                except Exception:
                    # Log and skip -- one bad message must not crash the bridge
                    logger.warning(
                        "translation failed (%s -> %s), skipping envelope",
                        source_protocol.value,
                        target_protocol.value,
                        exc_info=True,
                    )

    # ------------------------------------------------------------------
    # Translation dispatch
    # ------------------------------------------------------------------

    def _translate(self, envelope: Envelope, target_protocol: Protocol) -> Envelope:
        """Route to the appropriate translation method.

        If source and target protocols are the same, the envelope is
        returned with only the protocol tag updated (identity translation).

        Args:
            envelope: Incoming envelope to translate.
            target_protocol: Desired output protocol.

        Returns:
            A new :class:`Envelope` tagged with *target_protocol*.

        Example::

            translated = bridge._translate(envelope, Protocol.SLIM)
        """
        source_protocol = envelope.protocol

        if source_protocol == target_protocol:
            # Identity -- just retag so downstream sees the right protocol
            return Envelope(
                protocol=target_protocol,
                content_type=envelope.content_type,
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        key = (source_protocol, target_protocol)
        translator = self._translators.get(key)
        if translator is None:
            logger.warning(
                "no translator for %s -> %s, passing through as-is",
                source_protocol.value,
                target_protocol.value,
            )
            return Envelope(
                protocol=target_protocol,
                content_type=envelope.content_type,
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        return translator(envelope)

    # ------------------------------------------------------------------
    # A2A <-> SLIM
    # ------------------------------------------------------------------

    def _a2a_to_slim(self, envelope: Envelope) -> Envelope:
        """Wrap an A2A JSON payload inside a SLIM envelope.

        The A2A payload is preserved byte-for-byte; only the framing
        changes so SLIM consumers know the inner content type.

        Args:
            envelope: Envelope with ``Protocol.A2A`` and a JSON payload.

        Returns:
            SLIM-tagged envelope with ``application/a2a+json`` content type.

        Example::

            slim_env = bridge._a2a_to_slim(a2a_envelope)
            assert slim_env.content_type == "application/a2a+json"
        """
        # Validate that the payload is well-formed JSON before wrapping.
        # If it isn't, we still forward it but log a warning.
        try:
            json.loads(envelope.payload)
        except (json.JSONDecodeError, UnicodeDecodeError):
            logger.warning("A2A payload is not valid JSON, wrapping anyway")

        metadata = dict(envelope.metadata) if envelope.metadata else {}
        metadata["payload_type"] = "a2a"

        return Envelope(
            protocol=Protocol.SLIM,
            content_type="application/a2a+json",
            payload=envelope.payload,
            metadata=metadata,
        )

    def _slim_to_a2a(self, envelope: Envelope) -> Envelope:
        """Unwrap a SLIM envelope into an A2A message.

        If the SLIM content is already ``application/a2a+json`` the payload
        passes through unchanged.  Otherwise we synthesise a minimal A2A
        message wrapping the raw SLIM payload as a DataPart so that A2A
        consumers always receive a structurally valid message.

        Args:
            envelope: Envelope with ``Protocol.SLIM``.

        Returns:
            A2A-tagged envelope.

        Example::

            a2a_env = bridge._slim_to_a2a(slim_envelope)
            assert a2a_env.protocol == Protocol.A2A
        """
        if envelope.content_type == "application/a2a+json":
            # Already A2A JSON -- just retag
            return Envelope(
                protocol=Protocol.A2A,
                content_type="application/json",
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        # Wrap opaque SLIM payload in an A2A-shaped message so the receiver
        # can process it without special-casing raw bytes.
        try:
            # Attempt to decode as text for a friendlier representation
            payload_value = json.loads(envelope.payload)
        except (json.JSONDecodeError, UnicodeDecodeError):
            payload_value = base64.b64encode(envelope.payload).decode("ascii")

        a2a_message = {
            "role": "agent",
            "parts": [{"type": "data", "data": payload_value}],
        }
        return Envelope(
            protocol=Protocol.A2A,
            content_type="application/json",
            payload=json.dumps(a2a_message).encode("utf-8"),
            metadata=envelope.metadata,
        )

    # ------------------------------------------------------------------
    # MCP <-> SLIM
    # ------------------------------------------------------------------

    def _mcp_to_slim(self, envelope: Envelope) -> Envelope:
        """Wrap an MCP JSON-RPC payload inside a SLIM envelope.

        Extracts the JSON-RPC method name (if present) into metadata so
        that SLIM-side routers can filter without parsing the payload.

        Args:
            envelope: Envelope with ``Protocol.MCP`` and a JSON-RPC payload.

        Returns:
            SLIM-tagged envelope with ``application/mcp+json`` content type.

        Example::

            slim_env = bridge._mcp_to_slim(mcp_envelope)
            assert slim_env.content_type == "application/mcp+json"
        """
        metadata = dict(envelope.metadata) if envelope.metadata else {}
        metadata["payload_type"] = "mcp"

        # Pull the method name out for routing metadata
        try:
            rpc = json.loads(envelope.payload)
            method = rpc.get("method")
            if method:
                metadata["mcp_method"] = method
        except (json.JSONDecodeError, UnicodeDecodeError):
            logger.warning("MCP payload is not valid JSON, wrapping anyway")

        return Envelope(
            protocol=Protocol.SLIM,
            content_type="application/mcp+json",
            payload=envelope.payload,
            metadata=metadata,
        )

    def _slim_to_mcp(self, envelope: Envelope) -> Envelope:
        """Unwrap a SLIM envelope into an MCP JSON-RPC message.

        If the SLIM content is already ``application/mcp+json`` the payload
        passes through unchanged.  Otherwise we create a JSON-RPC
        ``message`` notification carrying the raw payload as base64 so
        that MCP consumers receive a spec-valid message.

        Args:
            envelope: Envelope with ``Protocol.SLIM``.

        Returns:
            MCP-tagged envelope.

        Example::

            mcp_env = bridge._slim_to_mcp(slim_envelope)
            assert mcp_env.protocol == Protocol.MCP
        """
        if envelope.content_type == "application/mcp+json":
            # Already MCP JSON-RPC -- just retag
            return Envelope(
                protocol=Protocol.MCP,
                content_type="application/json",
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        # Synthesise a JSON-RPC notification wrapping the opaque payload.
        # Notifications have no "id" field, so the receiver won't try to
        # send a response back.
        payload_b64 = base64.b64encode(envelope.payload).decode("ascii")
        notification = {
            "jsonrpc": "2.0",
            "method": "message",
            "params": {"data": payload_b64},
        }
        return Envelope(
            protocol=Protocol.MCP,
            content_type="application/json",
            payload=json.dumps(notification).encode("utf-8"),
            metadata=envelope.metadata,
        )

    # ------------------------------------------------------------------
    # A2A <-> MCP  (direct, skipping SLIM intermediary for efficiency)
    # ------------------------------------------------------------------

    def _a2a_to_mcp(self, envelope: Envelope) -> Envelope:
        """Translate an A2A message into an MCP JSON-RPC tool call request.

        Extracts text parts from the A2A message and packages them as
        arguments to an ``a2a_message`` tool call, preserving the A2A
        ``messageId`` as the JSON-RPC request id for correlation.

        Args:
            envelope: Envelope with ``Protocol.A2A`` and a JSON payload.

        Returns:
            MCP-tagged envelope containing a ``tools/call`` JSON-RPC request.

        Example::

            mcp_env = bridge._a2a_to_mcp(a2a_envelope)
            rpc = json.loads(mcp_env.payload)
            assert rpc["method"] == "tools/call"
        """
        try:
            a2a = json.loads(envelope.payload)
        except (json.JSONDecodeError, UnicodeDecodeError):
            # Can't parse -- wrap raw bytes as a data argument
            a2a = None

        if a2a is not None:
            parts = a2a.get("parts", [])
            msg_id = a2a.get("messageId", str(uuid.uuid4()))
        else:
            parts = [{"type": "data", "data": base64.b64encode(envelope.payload).decode("ascii")}]
            msg_id = str(uuid.uuid4())

        rpc_request = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "id": msg_id,
            "params": {
                "name": "a2a_message",
                "arguments": {"parts": parts},
            },
        }
        return Envelope(
            protocol=Protocol.MCP,
            content_type="application/json",
            payload=json.dumps(rpc_request).encode("utf-8"),
            metadata=envelope.metadata,
        )

    def _mcp_to_a2a(self, envelope: Envelope) -> Envelope:
        """Translate an MCP JSON-RPC message into an A2A message.

        Handles two cases:

        * **Result** (has ``"result"`` key): extracts content and wraps as
          A2A message parts (TextPart for strings, DataPart for everything
          else).
        * **Request** (has ``"method"`` key): wraps method + params as an
          A2A DataPart so the receiving agent can inspect the call.

        Args:
            envelope: Envelope with ``Protocol.MCP`` and a JSON-RPC payload.

        Returns:
            A2A-tagged envelope.

        Example::

            a2a_env = bridge._mcp_to_a2a(mcp_envelope)
            a2a = json.loads(a2a_env.payload)
            assert a2a["role"] == "agent"
        """
        try:
            rpc = json.loads(envelope.payload)
        except (json.JSONDecodeError, UnicodeDecodeError):
            rpc = None

        message_id = str(uuid.uuid4())
        parts = []

        if rpc is not None and "result" in rpc:
            # tools/call result -- extract content into A2A parts
            result = rpc["result"]
            content_items = result.get("content", []) if isinstance(result, dict) else []

            if not content_items and isinstance(result, dict):
                # Result is a plain dict without "content" -- wrap as data
                parts.append({"type": "data", "data": result})
            elif not content_items:
                # Scalar or non-dict result
                parts.append({"type": "text", "text": str(result)})
            else:
                for item in content_items:
                    if isinstance(item, dict) and item.get("type") == "text":
                        parts.append({"type": "text", "text": item.get("text", "")})
                    else:
                        parts.append({"type": "data", "data": item})

        elif rpc is not None and "method" in rpc:
            # Request or notification -- wrap the whole call as a data part
            # so the A2A agent can decide how to handle it
            parts.append({
                "type": "data",
                "data": {
                    "mcp_method": rpc["method"],
                    "mcp_params": rpc.get("params"),
                },
            })
        else:
            # Unrecognised structure -- pass raw
            raw = envelope.payload.decode("utf-8", errors="replace")
            parts.append({"type": "text", "text": raw})

        a2a_message = {
            "messageId": message_id,
            "role": "agent",
            "parts": parts,
        }
        return Envelope(
            protocol=Protocol.A2A,
            content_type="application/json",
            payload=json.dumps(a2a_message).encode("utf-8"),
            metadata=envelope.metadata,
        )

    # ------------------------------------------------------------------
    # ANP <-> SLIM
    # ------------------------------------------------------------------

    def _anp_to_slim(self, envelope: Envelope) -> Envelope:
        """Wrap an ANP JSON payload inside a SLIM envelope.

        Args:
            envelope: Envelope with ``Protocol.ANP`` and a JSON payload.

        Returns:
            SLIM-tagged envelope with ``application/anp+json`` content type.

        Example::

            slim_env = bridge._anp_to_slim(anp_envelope)
        """
        metadata = dict(envelope.metadata) if envelope.metadata else {}
        metadata["payload_type"] = "anp"

        return Envelope(
            protocol=Protocol.SLIM,
            content_type="application/anp+json",
            payload=envelope.payload,
            metadata=metadata,
        )

    def _slim_to_anp(self, envelope: Envelope) -> Envelope:
        """Unwrap a SLIM envelope into an ANP message.

        If the SLIM content is ``application/anp+json`` the payload passes
        through.  Otherwise wraps as a generic ANP message.

        Args:
            envelope: Envelope with ``Protocol.SLIM``.

        Returns:
            ANP-tagged envelope.

        Example::

            anp_env = bridge._slim_to_anp(slim_envelope)
        """
        if envelope.content_type == "application/anp+json":
            return Envelope(
                protocol=Protocol.ANP,
                content_type="application/json",
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        anp_wrapper = {
            "type": "bridged",
            "payload": json.loads(envelope.payload.decode("utf-8", errors="replace"))
            if envelope.payload
            else None,
        }
        return Envelope(
            protocol=Protocol.ANP,
            content_type="application/json",
            payload=json.dumps(anp_wrapper).encode("utf-8"),
            metadata=envelope.metadata,
        )

    # ------------------------------------------------------------------
    # LMOS <-> SLIM
    # ------------------------------------------------------------------

    def _lmos_to_slim(self, envelope: Envelope) -> Envelope:
        """Wrap an LMOS JSON payload inside a SLIM envelope.

        Args:
            envelope: Envelope with ``Protocol.LMOS`` and a JSON payload.

        Returns:
            SLIM-tagged envelope with ``application/lmos+json`` content type.

        Example::

            slim_env = bridge._lmos_to_slim(lmos_envelope)
        """
        metadata = dict(envelope.metadata) if envelope.metadata else {}
        metadata["payload_type"] = "lmos"

        return Envelope(
            protocol=Protocol.SLIM,
            content_type="application/lmos+json",
            payload=envelope.payload,
            metadata=metadata,
        )

    def _slim_to_lmos(self, envelope: Envelope) -> Envelope:
        """Unwrap a SLIM envelope into an LMOS message.

        If the SLIM content is ``application/lmos+json`` the payload passes
        through.  Otherwise wraps as a generic LMOS message.

        Args:
            envelope: Envelope with ``Protocol.SLIM``.

        Returns:
            LMOS-tagged envelope.

        Example::

            lmos_env = bridge._slim_to_lmos(slim_envelope)
        """
        if envelope.content_type == "application/lmos+json":
            return Envelope(
                protocol=Protocol.LMOS,
                content_type="application/json",
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        lmos_wrapper = {
            "type": "bridged",
            "payload": json.loads(envelope.payload.decode("utf-8", errors="replace"))
            if envelope.payload
            else None,
        }
        return Envelope(
            protocol=Protocol.LMOS,
            content_type="application/json",
            payload=json.dumps(lmos_wrapper).encode("utf-8"),
            metadata=envelope.metadata,
        )

    # ------------------------------------------------------------------
    # NLIP <-> SLIM
    # ------------------------------------------------------------------

    def _nlip_to_slim(self, envelope: Envelope) -> Envelope:
        """Wrap an NLIP JSON payload inside a SLIM envelope.

        Args:
            envelope: Envelope with ``Protocol.NLIP`` and a JSON payload.

        Returns:
            SLIM-tagged envelope with ``application/nlip+json`` content type.

        Example::

            slim_env = bridge._nlip_to_slim(nlip_envelope)
        """
        metadata = dict(envelope.metadata) if envelope.metadata else {}
        metadata["payload_type"] = "nlip"

        return Envelope(
            protocol=Protocol.SLIM,
            content_type="application/nlip+json",
            payload=envelope.payload,
            metadata=metadata,
        )

    def _slim_to_nlip(self, envelope: Envelope) -> Envelope:
        """Unwrap a SLIM envelope into an NLIP message.

        If the SLIM content is ``application/nlip+json`` the payload passes
        through.  Otherwise wraps as a generic NLIP message.

        Args:
            envelope: Envelope with ``Protocol.SLIM``.

        Returns:
            NLIP-tagged envelope.

        Example::

            nlip_env = bridge._slim_to_nlip(slim_envelope)
        """
        if envelope.content_type == "application/nlip+json":
            return Envelope(
                protocol=Protocol.NLIP,
                content_type="application/json",
                payload=envelope.payload,
                metadata=envelope.metadata,
            )

        nlip_wrapper = {
            "type": "bridged",
            "submessages": [{
                "subMessageType": "content",
                "contentType": "application/octet-stream",
                "content": json.loads(envelope.payload.decode("utf-8", errors="replace"))
                if envelope.payload
                else None,
            }],
        }
        return Envelope(
            protocol=Protocol.NLIP,
            content_type="application/json",
            payload=json.dumps(nlip_wrapper).encode("utf-8"),
            metadata=envelope.metadata,
        )
