"""Wire envelope schemas.

Four shapes — query / reply / ack / discovery — covering the lifecycle of an
agent-to-agent exchange:

- ``Query``: "I want this from you"
- ``Reply``: "Here's what you asked for, on the subject your query came in on"
- ``Ack``: "I saw your message and I'm alive"
- ``Discovery``: "Here's what I respond to" (sent on join + on demand)

All envelopes share ``from`` (sender identity) and ``ts`` (ISO-8601 UTC).
``type`` is a literal discriminator so a single ``parse_envelope`` can return
the right concrete model. ``from`` is the python-reserved-keyword issue we
side-step with pydantic's ``alias="from"`` on the field.
"""
from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, ValidationError


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


class _EnvelopeBase(BaseModel):
    model_config = ConfigDict(populate_by_name=True, extra="allow")

    type: str
    from_: str = Field(..., alias="from")
    ts: str = Field(default_factory=_now_iso)

    def to_json(self) -> str:
        return json.dumps(self.model_dump(by_alias=True, exclude_none=True), separators=(",", ":"))

    def to_bytes(self) -> bytes:
        return self.to_json().encode("utf-8")


class Query(_EnvelopeBase):
    type: Literal["query"] = "query"
    to: str
    query: str
    fields: list[str] | None = None
    reply_to: str | None = None
    request_id: str | None = None


class Reply(_EnvelopeBase):
    type: Literal["reply"] = "reply"
    reply_for: str
    in_reply_to: str | None = None
    request_id: str | None = None
    data: Any = None
    note: str | None = None


class Ack(_EnvelopeBase):
    type: Literal["ack"] = "ack"
    ack_for: str
    received_subjects: list[str] | None = None
    note: str | None = None
    alive: bool = True


class Discovery(_EnvelopeBase):
    type: Literal["discovery"] = "discovery"
    role: str
    cluster: str | None = None
    subjects_owned: list[str] = Field(default_factory=list)
    capabilities: dict[str, Any] = Field(default_factory=dict)


Envelope = Query | Reply | Ack | Discovery


_BY_TYPE: dict[str, type[_EnvelopeBase]] = {
    "query": Query,
    "reply": Reply,
    "ack": Ack,
    "discovery": Discovery,
}


class EnvelopeError(ValueError):
    """Raised when bytes / dict / str cannot be parsed into a known envelope."""


def parse_envelope(payload: bytes | str | dict[str, Any]) -> Envelope:
    """Decode a NATS message body into the right concrete envelope.

    Discriminates on ``type``. Unknown ``type`` values raise ``EnvelopeError``
    so callers can choose to log-and-drop or surface to the operator.
    """
    if isinstance(payload, (bytes, bytearray)):
        try:
            payload = payload.decode("utf-8")
        except UnicodeDecodeError as e:
            raise EnvelopeError(f"non-utf8 payload: {e}") from e
    if isinstance(payload, str):
        try:
            payload = json.loads(payload)
        except json.JSONDecodeError as e:
            raise EnvelopeError(f"invalid json: {e}") from e
    if not isinstance(payload, dict):
        raise EnvelopeError(f"envelope must decode to dict, got {type(payload).__name__}")

    type_ = payload.get("type")
    if type_ not in _BY_TYPE:
        raise EnvelopeError(f"unknown envelope type: {type_!r}")
    try:
        return _BY_TYPE[type_].model_validate(payload)
    except ValidationError as e:
        raise EnvelopeError(f"envelope validation failed for type={type_!r}: {e}") from e
