from __future__ import annotations

import math
import struct
from dataclasses import dataclass
from typing import NamedTuple

import lz4.block

from k1link.device_plugins.xgrids_k1.protocol.protobuf_wire import (
    ProtobufWireError,
    ProtoField,
    decode_zigzag64,
    iter_fields,
)


class StreamDecodeError(ValueError):
    """Raised when a K1 stream payload violates its verified bounds or schema."""


class UnsupportedCompressionError(StreamDecodeError):
    """Raised for a protocol compression type that has not been verified."""


@dataclass(frozen=True, slots=True)
class DecodeLimits:
    max_mqtt_payload_bytes: int = 2 * 1024 * 1024
    max_compressed_bytes: int = 1024 * 1024
    max_decompressed_bytes: int = 8 * 1024 * 1024
    max_compression_ratio: int = 64
    max_points_per_frame: int = 250_000

    def __post_init__(self) -> None:
        values = (
            self.max_mqtt_payload_bytes,
            self.max_compressed_bytes,
            self.max_decompressed_bytes,
            self.max_compression_ratio,
            self.max_points_per_frame,
        )
        if any(value < 1 for value in values):
            raise ValueError("all decode limits must be positive")


@dataclass(frozen=True, slots=True)
class MqttHeader:
    seq: int
    stamp: int
    scaler: int
    device_id: str
    session_id: str
    openapi_key: str | None


class LioPoint(NamedTuple):
    x_raw: int
    y_raw: int
    z_raw: int
    rgbi: int

    @property
    def intensity(self) -> int:
        """Return the only RGBA interpretation verified in the application."""
        return self.rgbi & 0xFF

    def scaled_xyz(self, scaler: int) -> tuple[float, float, float]:
        if scaler == 0:
            raise StreamDecodeError("point scaler is zero")
        return self.x_raw / scaler, self.y_raw / scaler, self.z_raw / scaler


@dataclass(frozen=True, slots=True)
class LioPointCloudFrame:
    header: MqttHeader
    compression: int
    compressed_bytes: int
    decompressed_bytes: int
    points: tuple[LioPoint, ...]


@dataclass(frozen=True, slots=True)
class LioPoseFrame:
    header: MqttHeader
    pose_stamp: int
    position_xyz: tuple[float, float, float]
    orientation_xyzw: tuple[float, float, float, float]
    distance: float
    pose_accuracy: float


class LegacyPoint(NamedTuple):
    x: float
    y: float
    z: float
    r: int
    g: int
    b: int
    intensity: int


@dataclass(frozen=True, slots=True)
class LegacyPointCloudFrame:
    envelope: bytes
    stride: int
    points: tuple[LegacyPoint, ...]


@dataclass(frozen=True, slots=True)
class LegacyPoseFrame:
    position_xyz: tuple[float, float, float]
    orientation_xyzw: tuple[float, float, float, float]
    skipped_offset_12: bytes
    unknown_tail: bytes


def _int_value(field: ProtoField, name: str) -> int:
    if field.wire_type != 0 or not isinstance(field.value, int):
        raise StreamDecodeError(f"{name} has the wrong protobuf wire type")
    return field.value


def _bytes_value(field: ProtoField, name: str, wire_type: int = 2) -> bytes:
    if field.wire_type != wire_type or not isinstance(field.value, bytes):
        raise StreamDecodeError(f"{name} has the wrong protobuf wire type")
    return field.value


def _float32(field: ProtoField, name: str) -> float:
    value = float(struct.unpack("<f", _bytes_value(field, name, 5))[0])
    if not math.isfinite(value):
        raise StreamDecodeError(f"{name} is not finite")
    return value


def _float64(field: ProtoField, name: str) -> float:
    value = float(struct.unpack("<d", _bytes_value(field, name, 1))[0])
    if not math.isfinite(value):
        raise StreamDecodeError(f"{name} is not finite")
    return value


def _text(field: ProtoField, name: str) -> str:
    try:
        return _bytes_value(field, name).decode("utf-8")
    except UnicodeDecodeError as exc:
        raise StreamDecodeError(f"{name} is not valid UTF-8") from exc


def _decode_header(payload: bytes) -> MqttHeader:
    seq = 0
    stamp = 0
    scaler = 0
    device_id = ""
    session_id = ""
    openapi_key: str | None = None
    for field in iter_fields(payload, max_fields=64):
        if field.number == 1:
            seq = _int_value(field, "header.seq")
        elif field.number == 2:
            stamp = decode_zigzag64(_int_value(field, "header.stamp"))
        elif field.number == 3:
            scaler = decode_zigzag64(_int_value(field, "header.scaler"))
        elif field.number == 4:
            device_id = _text(field, "header.device_id")
        elif field.number == 5:
            session_id = _text(field, "header.session_id")
        elif field.number == 6:
            openapi_key = _text(field, "header.openapi_key")
    return MqttHeader(seq, stamp, scaler, device_id, session_id, openapi_key)


def _decode_lio_point(payload: bytes) -> LioPoint:
    x_raw = 0
    y_raw = 0
    z_raw = 0
    rgbi = 0
    for field in iter_fields(payload, max_fields=16):
        if field.number == 1:
            x_raw = decode_zigzag64(_int_value(field, "point.x"))
        elif field.number == 2:
            y_raw = decode_zigzag64(_int_value(field, "point.y"))
        elif field.number == 3:
            z_raw = decode_zigzag64(_int_value(field, "point.z"))
        elif field.number == 4:
            rgbi = _int_value(field, "point.rgbi") & 0xFFFFFFFF
    return LioPoint(x_raw, y_raw, z_raw, rgbi)


def _decode_lio_pcl_report(
    payload: bytes,
    limits: DecodeLimits,
) -> tuple[MqttHeader, tuple[LioPoint, ...]]:
    header: MqttHeader | None = None
    points: list[LioPoint] = []
    try:
        for field in iter_fields(payload, max_fields=limits.max_points_per_frame + 64):
            if field.number == 1:
                header = _decode_header(_bytes_value(field, "lio_pcl.header"))
            elif field.number == 2:
                if len(points) >= limits.max_points_per_frame:
                    raise StreamDecodeError(
                        f"point frame exceeds {limits.max_points_per_frame} points"
                    )
                points.append(_decode_lio_point(_bytes_value(field, "lio_pcl.point")))
    except ProtobufWireError as exc:
        raise StreamDecodeError(f"invalid LioPclReport: {exc}") from exc

    if header is None:
        raise StreamDecodeError("LioPclReport has no header")
    if header.scaler == 0:
        raise StreamDecodeError("LioPclReport header scaler is zero")
    if not points:
        raise StreamDecodeError("LioPclReport has no points")
    return header, tuple(points)


def decode_lio_pcl(payload: bytes, limits: DecodeLimits | None = None) -> LioPointCloudFrame:
    """Decode the verified K1 lio_pcl envelope and raw LZ4 protobuf block."""
    bounds = limits or DecodeLimits()
    if len(payload) > bounds.max_mqtt_payload_bytes:
        raise StreamDecodeError("lio_pcl MQTT payload exceeds configured limit")

    compression = 0
    decompressed_size = 0
    compressed_data: bytes | None = None
    try:
        for field in iter_fields(payload, max_fields=64):
            if field.number == 2:
                compression = _int_value(field, "compression")
            elif field.number == 3:
                decompressed_size = _int_value(field, "compressed_size")
            elif field.number == 4:
                compressed_data = _bytes_value(field, "compressed_data")
    except ProtobufWireError as exc:
        raise StreamDecodeError(f"invalid MqttCompressMsg: {exc}") from exc

    if compression != 0:
        raise UnsupportedCompressionError(
            f"compression enum {compression} is not the verified raw-LZ4 mode"
        )
    if compressed_data is None or not compressed_data:
        raise StreamDecodeError("MqttCompressMsg has no compressed_data")
    if len(compressed_data) > bounds.max_compressed_bytes:
        raise StreamDecodeError("compressed_data exceeds configured limit")
    if decompressed_size < 1 or decompressed_size > bounds.max_decompressed_bytes:
        raise StreamDecodeError("compressed_size is outside configured bounds")
    if decompressed_size > len(compressed_data) * bounds.max_compression_ratio:
        raise StreamDecodeError("claimed LZ4 expansion ratio exceeds configured limit")

    try:
        decompressed = lz4.block.decompress(
            compressed_data,
            uncompressed_size=decompressed_size,
        )
    except lz4.block.LZ4BlockError as exc:
        raise StreamDecodeError(f"raw LZ4 decode failed: {exc}") from exc
    if len(decompressed) != decompressed_size:
        raise StreamDecodeError(
            f"raw LZ4 length mismatch: expected {decompressed_size}, got {len(decompressed)}"
        )

    header, points = _decode_lio_pcl_report(decompressed, bounds)
    return LioPointCloudFrame(
        header=header,
        compression=compression,
        compressed_bytes=len(compressed_data),
        decompressed_bytes=len(decompressed),
        points=points,
    )


def _decode_position(payload: bytes) -> tuple[float, float, float]:
    values = [0.0, 0.0, 0.0]
    for field in iter_fields(payload, max_fields=16):
        if 1 <= field.number <= 3:
            values[field.number - 1] = _float64(field, f"position.{field.number}")
    return values[0], values[1], values[2]


def _decode_orientation(payload: bytes) -> tuple[float, float, float, float]:
    values = [0.0, 0.0, 0.0, 0.0]
    for field in iter_fields(payload, max_fields=16):
        if 1 <= field.number <= 4:
            values[field.number - 1] = _float64(field, f"orientation.{field.number}")
    return values[0], values[1], values[2], values[3]


def _decode_pose(
    payload: bytes,
) -> tuple[tuple[float, float, float], tuple[float, float, float, float]]:
    position = (0.0, 0.0, 0.0)
    orientation = (0.0, 0.0, 0.0, 0.0)
    for field in iter_fields(payload, max_fields=16):
        if field.number == 1:
            position = _decode_position(_bytes_value(field, "pose.position"))
        elif field.number == 2:
            orientation = _decode_orientation(_bytes_value(field, "pose.orientation"))
    return position, orientation


def _decode_pose_stamped(
    payload: bytes,
) -> tuple[int, tuple[float, float, float], tuple[float, float, float, float]]:
    stamp = 0
    position = (0.0, 0.0, 0.0)
    orientation = (0.0, 0.0, 0.0, 0.0)
    for field in iter_fields(payload, max_fields=16):
        if field.number == 1:
            stamp = decode_zigzag64(_int_value(field, "pose_stamp.stamp"))
        elif field.number == 2:
            position, orientation = _decode_pose(_bytes_value(field, "pose_stamp.pose"))
    return stamp, position, orientation


def decode_lio_pose(payload: bytes, limits: DecodeLimits | None = None) -> LioPoseFrame:
    """Decode a direct lixel/application/report/lio_pose protobuf payload."""
    bounds = limits or DecodeLimits()
    if len(payload) > bounds.max_mqtt_payload_bytes:
        raise StreamDecodeError("lio_pose MQTT payload exceeds configured limit")

    header: MqttHeader | None = None
    pose_stamp = 0
    position = (0.0, 0.0, 0.0)
    orientation = (0.0, 0.0, 0.0, 0.0)
    distance = 0.0
    pose_accuracy = 0.0
    try:
        for field in iter_fields(payload, max_fields=64):
            if field.number == 1:
                header = _decode_header(_bytes_value(field, "lio_pose.header"))
            elif field.number == 2:
                pose_stamp, position, orientation = _decode_pose_stamped(
                    _bytes_value(field, "lio_pose.pose")
                )
            elif field.number == 3:
                distance = _float32(field, "lio_pose.distance")
            elif field.number == 4:
                pose_accuracy = _float32(field, "lio_pose.pose_accuracy")
    except ProtobufWireError as exc:
        raise StreamDecodeError(f"invalid LioPoseReport: {exc}") from exc

    if header is None:
        raise StreamDecodeError("LioPoseReport has no header")
    return LioPoseFrame(
        header=header,
        pose_stamp=pose_stamp,
        position_xyz=position,
        orientation_xyzw=orientation,
        distance=distance,
        pose_accuracy=pose_accuracy,
    )


def decode_legacy_pointcloud(
    payload: bytes,
    *,
    max_points: int = 250_000,
) -> LegacyPointCloudFrame:
    """Decode the verified legacy RealtimePointcloud envelope and point records."""
    if max_points < 1:
        raise ValueError("max_points must be positive")
    if len(payload) < 12:
        raise StreamDecodeError("legacy pointcloud payload is shorter than 12-byte envelope")
    stride = int.from_bytes(payload[:4], "little")
    if stride < 15:
        raise StreamDecodeError("legacy point stride is smaller than xyz+rgb")
    body = payload[12:]
    if len(body) % stride:
        raise StreamDecodeError("legacy pointcloud body is not divisible by stride")
    point_count = len(body) // stride
    if point_count > max_points:
        raise StreamDecodeError(f"legacy pointcloud exceeds {max_points} points")

    points: list[LegacyPoint] = []
    for offset in range(0, len(body), stride):
        x, y, z = struct.unpack_from("<fff", body, offset)
        if not all(math.isfinite(value) for value in (x, y, z)):
            raise StreamDecodeError("legacy point position is not finite")
        r, g, b = body[offset + 12 : offset + 15]
        intensity = body[offset + 15] if stride >= 16 else 255
        points.append(LegacyPoint(x, y, z, r, g, b, intensity))
    return LegacyPointCloudFrame(payload[:12], stride, tuple(points))


def decode_legacy_pose(payload: bytes) -> LegacyPoseFrame:
    """Decode the verified legacy RealtimePath position/quaternion fields."""
    if len(payload) < 32:
        raise StreamDecodeError("legacy pose payload is shorter than 32 bytes")
    x, y, z = struct.unpack_from("<fff", payload, 0)
    wire_w, qx, qy, qz = struct.unpack_from("<ffff", payload, 16)
    values = (x, y, z, qx, qy, qz, wire_w)
    if not all(math.isfinite(value) for value in values):
        raise StreamDecodeError("legacy pose contains a non-finite value")
    return LegacyPoseFrame(
        position_xyz=(x, y, z),
        orientation_xyzw=(qx, qy, qz, wire_w),
        skipped_offset_12=payload[12:16],
        unknown_tail=payload[32:],
    )


def decode_pre_path_array(payload: bytes) -> tuple[float, ...]:
    """Decode the exact 16-float64 legacy PrePathArray matrix."""
    if len(payload) != 128:
        raise StreamDecodeError("PrePathArray payload must be exactly 128 bytes")
    values = struct.unpack("<16d", payload)
    if not all(math.isfinite(value) for value in values):
        raise StreamDecodeError("PrePathArray contains a non-finite value")
    return values
