from __future__ import annotations

from typing import Protocol

from k1link.data_plane import (
    ConsumerFrameContext,
    DecodedDataPlaneView,
    DecodedPointCloudView,
    DecodedPoseView,
    NormalizationError,
)
from k1link.device_plugins.xgrids_k1.protocol.streams import (
    LegacyPointCloudFrame,
    LegacyPoseFrame,
    LioPointCloudFrame,
    LioPoseFrame,
    StreamDecodeError,
    decode_legacy_pointcloud,
    decode_legacy_pose,
    decode_lio_pcl,
    decode_lio_pose,
)

CANONICAL_MAP_FRAME = "map"
CANONICAL_SENSOR_FRAME = "sensor"


class K1TransportMessage(Protocol):
    """Minimum raw transport shape accepted by the K1 normalizer."""

    @property
    def sequence(self) -> int: ...

    @property
    def topic(self) -> str: ...

    @property
    def payload(self) -> bytes: ...

    @property
    def received_at_epoch_ns(self) -> int: ...

    @property
    def received_monotonic_ns(self) -> int | None: ...

    @property
    def source(self) -> str: ...


def normalize_k1_message(
    message: K1TransportMessage,
    *,
    processing_started_monotonic_ns: int,
) -> DecodedDataPlaneView | None:
    """Normalize one verified K1 transport message.

    Unknown channels intentionally return ``None``. A known channel with an
    invalid payload raises a transport-neutral ``NormalizationError``. Neither
    raw channel names nor raw bytes are retained in the returned consumer view.
    """

    try:
        if message.topic.endswith("/lio_pcl"):
            return _normalize_lio_points(
                decode_lio_pcl(message.payload),
                _context(message, processing_started_monotonic_ns),
            )
        if message.topic == "RealtimePointcloud":
            return _normalize_legacy_points(
                decode_legacy_pointcloud(message.payload),
                _context(message, processing_started_monotonic_ns),
            )
        if message.topic.endswith("/lio_pose"):
            return _normalize_lio_pose(
                decode_lio_pose(message.payload),
                _context(message, processing_started_monotonic_ns),
            )
        if message.topic == "RealtimePath":
            return _normalize_legacy_pose(
                decode_legacy_pose(message.payload),
                _context(message, processing_started_monotonic_ns),
            )
    except StreamDecodeError as exc:
        raise NormalizationError("known K1 data-plane message failed validation") from exc
    return None


def _context(
    message: K1TransportMessage,
    processing_started_monotonic_ns: int,
) -> ConsumerFrameContext:
    return ConsumerFrameContext(
        sequence=message.sequence,
        captured_at_epoch_ns=message.received_at_epoch_ns,
        received_monotonic_ns=message.received_monotonic_ns,
        processing_started_monotonic_ns=processing_started_monotonic_ns,
        encoded_size_bytes=len(message.payload),
        live=message.source == "live_mqtt",
    )


def _with_device_context(
    context: ConsumerFrameContext,
    *,
    source_device_alias: str,
    source_session_alias: str,
) -> ConsumerFrameContext:
    return ConsumerFrameContext(
        sequence=context.sequence,
        captured_at_epoch_ns=context.captured_at_epoch_ns,
        received_monotonic_ns=context.received_monotonic_ns,
        processing_started_monotonic_ns=context.processing_started_monotonic_ns,
        encoded_size_bytes=context.encoded_size_bytes,
        live=context.live,
        source_device_alias=source_device_alias or None,
        source_session_alias=source_session_alias or None,
    )


def _normalize_lio_points(
    frame: LioPointCloudFrame,
    context: ConsumerFrameContext,
) -> DecodedPointCloudView:
    context = _with_device_context(
        context,
        source_device_alias=frame.header.device_id,
        source_session_alias=frame.header.session_id,
    )
    positions = tuple(point.scaled_xyz(frame.header.scaler) for point in frame.points)
    intensities = bytes(point.intensity for point in frame.points)
    return DecodedPointCloudView(
        context=context,
        frame_id=CANONICAL_MAP_FRAME,
        positions_xyz=positions,
        intensities=intensities,
    )


def _normalize_legacy_points(
    frame: LegacyPointCloudFrame,
    context: ConsumerFrameContext,
) -> DecodedPointCloudView:
    return DecodedPointCloudView(
        context=context,
        frame_id=CANONICAL_MAP_FRAME,
        positions_xyz=tuple((point.x, point.y, point.z) for point in frame.points),
        intensities=bytes(point.intensity for point in frame.points),
        colors_rgb=bytes(
            channel for point in frame.points for channel in (point.r, point.g, point.b)
        ),
    )


def _normalize_lio_pose(
    frame: LioPoseFrame,
    context: ConsumerFrameContext,
) -> DecodedPoseView:
    context = _with_device_context(
        context,
        source_device_alias=frame.header.device_id,
        source_session_alias=frame.header.session_id,
    )
    return DecodedPoseView(
        context=context,
        frame_id=CANONICAL_MAP_FRAME,
        child_frame_id=CANONICAL_SENSOR_FRAME,
        position_xyz=frame.position_xyz,
        orientation_xyzw=frame.orientation_xyzw,
    )


def _normalize_legacy_pose(
    frame: LegacyPoseFrame,
    context: ConsumerFrameContext,
) -> DecodedPoseView:
    return DecodedPoseView(
        context=context,
        frame_id=CANONICAL_MAP_FRAME,
        child_frame_id=CANONICAL_SENSOR_FRAME,
        position_xyz=frame.position_xyz,
        orientation_xyzw=frame.orientation_xyzw,
    )
