from __future__ import annotations

import math
import struct

import lz4.block
import pytest

from k1link.protocol.protobuf_wire import ProtobufWireError, decode_zigzag64, iter_fields
from k1link.protocol.streams import (
    DecodeLimits,
    StreamDecodeError,
    UnsupportedCompressionError,
    decode_legacy_pointcloud,
    decode_legacy_pose,
    decode_lio_pcl,
    decode_lio_pose,
    decode_pre_path_array,
)


def _varint(value: int) -> bytes:
    encoded = bytearray()
    while value > 0x7F:
        encoded.append((value & 0x7F) | 0x80)
        value >>= 7
    encoded.append(value)
    return bytes(encoded)


def _key(number: int, wire_type: int) -> bytes:
    return _varint((number << 3) | wire_type)


def _uint(number: int, value: int) -> bytes:
    return _key(number, 0) + _varint(value)


def _sint(number: int, value: int) -> bytes:
    zigzag = (value << 1) ^ (value >> 63)
    return _uint(number, zigzag & 0xFFFFFFFFFFFFFFFF)


def _bytes(number: int, value: bytes) -> bytes:
    return _key(number, 2) + _varint(len(value)) + value


def _fixed32(number: int, value: float) -> bytes:
    return _key(number, 5) + struct.pack("<f", value)


def _fixed64(number: int, value: float) -> bytes:
    return _key(number, 1) + struct.pack("<d", value)


def _header(*, scaler: int = 1000) -> bytes:
    return b"".join(
        (
            _uint(1, 7),
            _sint(2, 123456),
            _sint(3, scaler),
            _bytes(4, b"device-redacted"),
            _bytes(5, b"session-redacted"),
        )
    )


def _pcl_payload(*, compression: int = 0, scaler: int = 1000) -> bytes:
    point_1 = _sint(1, 1000) + _sint(2, -2000) + _sint(3, 500) + _uint(4, 0x11223344)
    point_2 = _sint(1, -250) + _sint(2, 0) + _sint(3, 4000) + _uint(4, 0xAABBCC09)
    report = _bytes(1, _header(scaler=scaler)) + _bytes(2, point_1) + _bytes(2, point_2)
    compressed = lz4.block.compress(report, store_size=False)
    fields = []
    if compression:
        fields.append(_uint(2, compression))
    fields.extend((_uint(3, len(report)), _bytes(4, compressed)))
    return b"".join(fields)


def test_protobuf_wire_zigzag_and_bounds() -> None:
    assert decode_zigzag64(0) == 0
    assert decode_zigzag64(1) == -1
    assert decode_zigzag64(2) == 1
    with pytest.raises(ProtobufWireError, match="truncated"):
        list(iter_fields(b"\x0a\x02\x01"))
    with pytest.raises(ProtobufWireError, match="unsupported"):
        list(iter_fields(b"\x0b"))


def test_decode_lio_pcl_raw_lz4() -> None:
    frame = decode_lio_pcl(_pcl_payload())
    assert frame.header.seq == 7
    assert frame.header.stamp == 123456
    assert frame.header.scaler == 1000
    assert len(frame.points) == 2
    assert frame.points[0].scaled_xyz(frame.header.scaler) == (1.0, -2.0, 0.5)
    assert frame.points[0].rgbi == 0x11223344
    assert frame.points[0].intensity == 0x44
    assert frame.points[1].scaled_xyz(frame.header.scaler) == (-0.25, 0.0, 4.0)
    assert frame.points[1].intensity == 9


def test_decode_lio_pcl_rejects_unverified_or_unsafe_frames() -> None:
    with pytest.raises(UnsupportedCompressionError, match="enum 1"):
        decode_lio_pcl(_pcl_payload(compression=1))
    with pytest.raises(StreamDecodeError, match="scaler is zero"):
        decode_lio_pcl(_pcl_payload(scaler=0))
    with pytest.raises(StreamDecodeError, match="exceeds 1 points"):
        decode_lio_pcl(_pcl_payload(), DecodeLimits(max_points_per_frame=1))
    with pytest.raises(StreamDecodeError, match="MQTT payload exceeds"):
        decode_lio_pcl(_pcl_payload(), DecodeLimits(max_mqtt_payload_bytes=4))


def test_decode_lio_pose() -> None:
    position = _fixed64(1, 1.25) + _fixed64(2, -2.5) + _fixed64(3, 3.75)
    orientation = (
        _fixed64(1, 0.1)
        + _fixed64(2, 0.2)
        + _fixed64(3, 0.3)
        + _fixed64(4, 0.9)
    )
    pose = _bytes(1, position) + _bytes(2, orientation)
    stamped = _sint(1, 987654321) + _bytes(2, pose)
    payload = (
        _bytes(1, _header())
        + _bytes(2, stamped)
        + _fixed32(3, 12.5)
        + _fixed32(4, 0.001)
    )

    frame = decode_lio_pose(payload)
    assert frame.pose_stamp == 987654321
    assert frame.position_xyz == (1.25, -2.5, 3.75)
    assert frame.orientation_xyzw == pytest.approx((0.1, 0.2, 0.3, 0.9))
    assert frame.distance == 12.5
    assert frame.pose_accuracy == pytest.approx(0.001)


def test_decode_lio_pose_rejects_nonfinite_float() -> None:
    payload = _bytes(1, _header()) + _fixed32(3, math.nan)
    with pytest.raises(StreamDecodeError, match="not finite"):
        decode_lio_pose(payload)


def test_decode_legacy_pointcloud() -> None:
    envelope = struct.pack("<III", 16, 123, 456)
    body = struct.pack("<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40)
    frame = decode_legacy_pointcloud(envelope + body)
    assert frame.stride == 16
    assert frame.envelope == envelope
    assert frame.points[0] == (1.0, -2.0, 3.0, 10, 20, 30, 40)


def test_decode_legacy_pose_and_matrix() -> None:
    payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
    frame = decode_legacy_pose(payload + b"tail")
    assert frame.position_xyz == (1.0, 2.0, 3.0)
    assert frame.orientation_xyzw == pytest.approx((0.1, 0.2, 0.3, 0.9))
    assert frame.skipped_offset_12 == struct.pack("<f", 99.0)
    assert frame.unknown_tail == b"tail"

    matrix = tuple(float(index) for index in range(16))
    assert decode_pre_path_array(struct.pack("<16d", *matrix)) == matrix
    with pytest.raises(StreamDecodeError, match="exactly 128"):
        decode_pre_path_array(b"short")