"""Msgpack codec for A2A envelopes."""

from __future__ import annotations

import msgpack

from p2p.errors import ProtocolError
from p2p.protocol.envelope import A2AEnvelope


def encode(envelope: A2AEnvelope) -> bytes:
    return msgpack.packb(envelope.model_dump(mode="python"), use_bin_type=True)


_MAX_MESSAGE_SIZE = 10 * 1024 * 1024  # 10 MB


def decode(raw: bytes, *, max_size: int = _MAX_MESSAGE_SIZE) -> A2AEnvelope:
    if len(raw) > max_size:
        raise ProtocolError(f"Message too large: {len(raw)} > {max_size}")
    try:
        data = msgpack.unpackb(raw, raw=False)
    except Exception as e:
        raise ProtocolError(f"Bad msgpack: {e}") from e
    if not isinstance(data, dict):
        raise ProtocolError(f"Expected dict, got {type(data).__name__}")
    try:
        return A2AEnvelope.model_validate(data)
    except Exception as e:
        raise ProtocolError(f"Invalid envelope: {e}") from e
