"""Transport-neutral messages used inside an edge-ai pipeline.

The envelope is deliberately independent from protobuf, HTTP and model-specific
types.  Connectors translate their wire representation at the edge of the
pipeline and operators only exchange :class:`Envelope` instances.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from time import time_ns
from types import MappingProxyType
from typing import Any, Generic, Mapping, TypeVar
from uuid import uuid4


PayloadT = TypeVar("PayloadT")
NewPayloadT = TypeVar("NewPayloadT")


def _new_trace_id() -> str:
    return uuid4().hex


@dataclass(frozen=True, slots=True)
class Envelope(Generic[PayloadT]):
    """An immutable metadata wrapper around one pipeline payload.

    ``Envelope`` immutability is shallow: the wrapper and its attributes cannot
    be changed, but callers remain responsible for treating a mutable payload
    (for example a NumPy array) as read-only after publishing it.  This permits
    zero-copy fan-out without imposing a dependency on a particular array or
    shared-memory implementation.

    All timestamps use Unix time in nanoseconds so timestamps received from
    remote processes can be compared.  ``deadline_ns`` is optional; runtimes do
    not drop expired messages automatically because expiry policy is a business
    decision best implemented by an operator or sink.
    """

    payload: PayloadT
    schema_name: str = "opaque"
    schema_version: int = 1
    source_id: str = ""
    sequence: int = 0
    captured_at_ns: int | None = None
    received_at_ns: int = field(default_factory=time_ns)
    deadline_ns: int | None = None
    trace_id: str = field(default_factory=_new_trace_id)
    session_id: str | None = None
    attributes: Mapping[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.schema_name:
            raise ValueError("schema_name must not be empty")
        if self.schema_version < 1:
            raise ValueError("schema_version must be at least 1")
        if self.sequence < 0:
            raise ValueError("sequence must not be negative")
        if not self.trace_id:
            raise ValueError("trace_id must not be empty")
        if self.received_at_ns < 0:
            raise ValueError("received_at_ns must not be negative")
        if self.captured_at_ns is not None and self.captured_at_ns < 0:
            raise ValueError("captured_at_ns must not be negative")
        if self.deadline_ns is not None and self.deadline_ns < 0:
            raise ValueError("deadline_ns must not be negative")

        copied_attributes = dict(self.attributes)
        if any(not isinstance(key, str) for key in copied_attributes):
            raise TypeError("attribute keys must be strings")
        object.__setattr__(
            self,
            "attributes",
            MappingProxyType(copied_attributes),
        )

    @property
    def schema(self) -> str:
        """Return the compact schema identifier, for example ``ImageFrame/v1``."""

        return f"{self.schema_name}/v{self.schema_version}"

    def is_expired(self, now_ns: int | None = None) -> bool:
        """Return whether the optional deadline has passed."""

        if self.deadline_ns is None:
            return False
        return (time_ns() if now_ns is None else now_ns) >= self.deadline_ns

    def with_payload(
        self,
        payload: NewPayloadT,
        *,
        schema_name: str | None = None,
        schema_version: int | None = None,
        source_id: str | None = None,
    ) -> Envelope[NewPayloadT]:
        """Create a derived envelope while retaining correlation metadata."""

        return Envelope(
            payload=payload,
            schema_name=self.schema_name if schema_name is None else schema_name,
            schema_version=(
                self.schema_version if schema_version is None else schema_version
            ),
            source_id=self.source_id if source_id is None else source_id,
            sequence=self.sequence,
            captured_at_ns=self.captured_at_ns,
            # Keep the ingress timestamp so end-to-end latency remains
            # observable across an arbitrary number of operators.
            received_at_ns=self.received_at_ns,
            deadline_ns=self.deadline_ns,
            trace_id=self.trace_id,
            session_id=self.session_id,
            attributes=self.attributes,
        )

    def with_attributes(self, **attributes: Any) -> Envelope[PayloadT]:
        """Create a copy with additional or replaced read-only attributes."""

        merged = dict(self.attributes)
        merged.update(attributes)
        return Envelope(
            payload=self.payload,
            schema_name=self.schema_name,
            schema_version=self.schema_version,
            source_id=self.source_id,
            sequence=self.sequence,
            captured_at_ns=self.captured_at_ns,
            received_at_ns=self.received_at_ns,
            deadline_ns=self.deadline_ns,
            trace_id=self.trace_id,
            session_id=self.session_id,
            attributes=merged,
        )


__all__ = ["Envelope", "PayloadT"]
