"""Vendor-neutral canonical stream envelopes."""

from __future__ import annotations

import math
from enum import StrEnum
from typing import Annotated, Literal

from pydantic import AwareDatetime, Field, NonNegativeInt, PositiveInt, model_validator

from .common import API_VERSION, ApiVersion, ContractModel, Identifier, JsonObject, ShortText
from .payloads import PayloadHandle
from .session import DeviceSessionRef


class ClockDomain(StrEnum):
    UTC = "utc"
    DEVICE_MONOTONIC = "device-monotonic"
    HOST_MONOTONIC = "host-monotonic"
    UNKNOWN = "unknown"


class CanonicalStreamHeader(ContractModel):
    api_version: ApiVersion = API_VERSION
    message_id: Identifier
    session: DeviceSessionRef
    acquisition_id: Identifier | None = None
    source_id: Identifier
    channel_id: Identifier
    sequence: NonNegativeInt
    captured_at: AwareDatetime
    observed_at: AwareDatetime
    clock_domain: ClockDomain
    compatibility_profile_id: Identifier | None = None


class PointFieldType(StrEnum):
    INT8 = "int8"
    UINT8 = "uint8"
    INT16 = "int16"
    UINT16 = "uint16"
    INT32 = "int32"
    UINT32 = "uint32"
    FLOAT32 = "float32"
    FLOAT64 = "float64"


_POINT_FIELD_BYTES = {
    PointFieldType.INT8: 1,
    PointFieldType.UINT8: 1,
    PointFieldType.INT16: 2,
    PointFieldType.UINT16: 2,
    PointFieldType.INT32: 4,
    PointFieldType.UINT32: 4,
    PointFieldType.FLOAT32: 4,
    PointFieldType.FLOAT64: 8,
}


class PointField(ContractModel):
    name: Identifier
    offset: NonNegativeInt
    data_type: PointFieldType
    count: PositiveInt = 1


class PointCloudFrame(ContractModel):
    kind: Literal["point-cloud"] = "point-cloud"
    header: CanonicalStreamHeader
    coordinate_frame: Identifier
    width: PositiveInt
    height: PositiveInt = 1
    point_count: PositiveInt
    point_step: PositiveInt
    row_step: PositiveInt
    little_endian: bool = True
    dense: bool
    fields: tuple[PointField, ...] = Field(min_length=3)
    payload: PayloadHandle

    @model_validator(mode="after")
    def point_layout_must_be_self_consistent(self) -> PointCloudFrame:
        if self.point_count != self.width * self.height:
            raise ValueError("point_count must equal width * height")
        if self.row_step != self.point_step * self.width:
            raise ValueError("row_step must equal point_step * width")
        if self.payload.byte_length != self.row_step * self.height:
            raise ValueError("payload byte_length must match the declared point layout")
        names = [field.name for field in self.fields]
        if len(names) != len(set(names)):
            raise ValueError("point field names must be unique")
        if not {"x", "y", "z"}.issubset(names):
            raise ValueError("canonical point clouds require x, y, and z fields")
        for field in self.fields:
            field_end = field.offset + _POINT_FIELD_BYTES[field.data_type] * field.count
            if field_end > self.point_step:
                raise ValueError(f"point field {field.name} exceeds point_step")
        return self


class Vector3(ContractModel):
    x: float
    y: float
    z: float


class Quaternion(ContractModel):
    x: float
    y: float
    z: float
    w: float

    @model_validator(mode="after")
    def quaternion_must_be_normalized(self) -> Quaternion:
        norm = math.sqrt(self.x**2 + self.y**2 + self.z**2 + self.w**2)
        if not math.isclose(norm, 1.0, rel_tol=1e-3, abs_tol=1e-3):
            raise ValueError("quaternion must be normalized")
        return self


class PoseFrame(ContractModel):
    kind: Literal["pose"] = "pose"
    header: CanonicalStreamHeader
    parent_frame: Identifier
    child_frame: Identifier
    translation_m: Vector3
    rotation: Quaternion
    covariance: tuple[float, ...] | None = None

    @model_validator(mode="after")
    def covariance_must_be_six_by_six(self) -> PoseFrame:
        if self.covariance is not None and len(self.covariance) != 36:
            raise ValueError("pose covariance must contain 36 values")
        return self


class ImageFrame(ContractModel):
    kind: Literal["image"] = "image"
    header: CanonicalStreamHeader
    camera_id: Identifier
    optical_frame: Identifier
    width: PositiveInt
    height: PositiveInt
    pixel_format: Literal["rgb8", "bgr8", "rgba8", "mono8", "mono16", "depth16", "jpeg", "png"]
    payload: PayloadHandle
    calibration_id: Identifier | None = None


class EncodedVideoPacket(ContractModel):
    kind: Literal["encoded-video"] = "encoded-video"
    header: CanonicalStreamHeader
    stream_id: Identifier
    camera_id: Identifier
    codec: Literal["h264", "h265", "mjpeg", "av1", "unknown"]
    key_frame: bool
    presentation_time_ns: NonNegativeInt
    decode_time_ns: NonNegativeInt | None = None
    duration_ns: NonNegativeInt | None = None
    payload: PayloadHandle


class DeviceStatusFrame(ContractModel):
    kind: Literal["device-status"] = "device-status"
    header: CanonicalStreamHeader
    health: Literal["unknown", "healthy", "degraded", "fault"]
    state: Identifier
    measurements: JsonObject = Field(default_factory=dict)
    message: ShortText | None = None


CanonicalStreamEnvelope = Annotated[
    PointCloudFrame | PoseFrame | ImageFrame | EncodedVideoPacket | DeviceStatusFrame,
    Field(discriminator="kind"),
]
