"""Message envelope for multi-protocol support.

Wraps payloads with protocol identification so receivers can deserialize
correctly when different protocols share a channel or bridge.

Wire format: ``[header_len:4 LE][json header][payload]``

The header is a compact JSON object::

    {"p": "a2a", "ct": "application/json", "md": {"key": "val"}}

Usage::

    from skytale_sdk.envelope import Envelope, Protocol

    env = Envelope(Protocol.A2A, "application/json", b'{"parts":[]}')
    data = env.serialize()
    env2 = Envelope.deserialize(data)
"""

from __future__ import annotations

import json
import struct
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional


class Protocol(str, Enum):
    """Supported wire protocols.

    Example::

        >>> Protocol.A2A.value
        'a2a'
    """

    RAW = "raw"
    A2A = "a2a"
    MCP = "mcp"
    SLIM = "slim"
    ACP = "acp"
    ANP = "anp"
    LMOS = "lmos"
    NLIP = "nlip"


# 4-byte little-endian unsigned int prefix
_HEADER_FMT = "<I"
_HEADER_SIZE = struct.calcsize(_HEADER_FMT)

# Minimum valid envelope: 4-byte header length + 2-byte JSON "{}" + 0-byte payload
_MIN_ENVELOPE_SIZE = _HEADER_SIZE + 2

# Maximum allowed header length to prevent memory exhaustion from crafted inputs
_MAX_HEADER_LEN = 65536


@dataclass(frozen=True)
class Envelope:
    """Protocol-tagged message envelope.

    Args:
        protocol: Source protocol identifier.
        content_type: MIME type of *payload* (e.g. ``"application/json"``).
        payload: Raw payload bytes.
        metadata: Optional key-value metadata. Values must be JSON-serializable.

    Example::

        >>> env = Envelope(Protocol.A2A, "application/json", b'{"hello":1}')
        >>> Envelope.deserialize(env.serialize()) == env
        True
    """

    protocol: Protocol
    content_type: str
    payload: bytes
    metadata: Optional[dict] = field(default=None)

    def serialize(self) -> bytes:
        """Serialize to wire format: ``[header_len:4 LE][json header][payload]``.

        Returns:
            Bytes ready for transmission through a Skytale channel.

        Example::

            >>> env = Envelope(Protocol.RAW, "text/plain", b"hello")
            >>> data = env.serialize()
            >>> len(data) > 4 + 5  # header prefix + payload
            True
        """
        header = {"p": self.protocol.value, "ct": self.content_type}
        if self.metadata:
            header["md"] = self.metadata
        header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
        return struct.pack(_HEADER_FMT, len(header_bytes)) + header_bytes + self.payload

    @staticmethod
    def deserialize(data: bytes) -> Envelope:
        """Deserialize from wire format.

        Args:
            data: Bytes produced by :meth:`serialize`.

        Returns:
            Reconstructed :class:`Envelope`.

        Raises:
            ValueError: If *data* is too short, the header length is invalid,
                or the header JSON is malformed.

        Example::

            >>> env = Envelope(Protocol.MCP, "application/json", b'{}')
            >>> Envelope.deserialize(env.serialize()).protocol
            <Protocol.MCP: 'mcp'>
        """
        if len(data) < _MIN_ENVELOPE_SIZE:
            raise ValueError(
                f"envelope too short: {len(data)} bytes "
                f"(minimum {_MIN_ENVELOPE_SIZE})"
            )

        (header_len,) = struct.unpack(_HEADER_FMT, data[:_HEADER_SIZE])

        if header_len > _MAX_HEADER_LEN:
            raise ValueError(
                f"header length {header_len} exceeds maximum "
                f"({_MAX_HEADER_LEN} bytes)"
            )

        if header_len > len(data) - _HEADER_SIZE:
            raise ValueError(
                f"header length {header_len} exceeds available data "
                f"({len(data) - _HEADER_SIZE} bytes)"
            )

        header_bytes = data[_HEADER_SIZE : _HEADER_SIZE + header_len]
        try:
            header = json.loads(header_bytes)
        except json.JSONDecodeError as exc:
            raise ValueError(f"malformed envelope header: {exc}") from exc

        try:
            protocol = Protocol(header["p"])
        except (KeyError, ValueError) as exc:
            raise ValueError(f"invalid protocol in envelope header: {exc}") from exc

        content_type = header.get("ct", "application/octet-stream")
        metadata = header.get("md")
        payload = data[_HEADER_SIZE + header_len :]

        return Envelope(
            protocol=protocol,
            content_type=content_type,
            payload=payload,
            metadata=metadata,
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Envelope):
            return NotImplemented
        return (
            self.protocol == other.protocol
            and self.content_type == other.content_type
            and self.payload == other.payload
            and self.metadata == other.metadata
        )
