from __future__ import annotations

import math
import struct
from types import SimpleNamespace

import lz4.block
import pytest

from k1link.device_plugins.xgrids_k1.protocol.modeling import (
    MODELING_REPORT_TOPIC,
    ModelingReportDecodeError,
    decode_modeling_report,
    is_modeling_report_topic,
    observe_modeling_report,
)
from k1link.device_plugins.xgrids_k1.protocol.protobuf_wire import (
    ProtobufWireError,
    decode_zigzag64,
    encode_zigzag64,
    iter_fields,
)
from k1link.device_plugins.xgrids_k1.protocol.streams import (
    DecodeLimits,
    StreamDecodeError,
    UnsupportedCompressionError,
    decode_legacy_pointcloud,
    decode_legacy_pose,
    decode_lio_pcl,
    decode_lio_pose,
    decode_pre_path_array,
)
from k1link.viewer.metrics import BridgeMetrics


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
    assert encode_zigzag64(0) == 0
    assert encode_zigzag64(-1) == 1
    assert encode_zigzag64(1) == 2
    assert decode_zigzag64(encode_zigzag64(-(1 << 63))) == -(1 << 63)
    assert decode_zigzag64(encode_zigzag64((1 << 63) - 1)) == (1 << 63) - 1
    with pytest.raises(ProtobufWireError, match="outside int64"):
        encode_zigzag64(1 << 63)
    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_modeling_report_device_telemetry() -> None:
    scan_status = _fixed32(1, 41.521404) + _fixed32(2, 1.375) + _uint(3, 204)
    payload = _bytes(1, _header()) + _uint(2, 73) + _bytes(3, scan_status)

    telemetry = decode_modeling_report(payload)

    assert telemetry.move_distance_meters == pytest.approx(41.521404)
    assert telemetry.move_speed_meters_per_second == pytest.approx(1.375)
    assert telemetry.scan_time_ticks == 204
    assert telemetry.elapsed_seconds == pytest.approx(102.0)
    assert telemetry.pgo_progress == 73
    assert is_modeling_report_topic(MODELING_REPORT_TOPIC)
    assert not is_modeling_report_topic(f"{MODELING_REPORT_TOPIC}/extra")

    metrics = BridgeMetrics()
    assert observe_modeling_report(
        SimpleNamespace(topic=MODELING_REPORT_TOPIC, payload=payload),
        metrics,
    )
    snapshot = metrics.snapshot()
    assert snapshot["device_elapsed_seconds"] == pytest.approx(102.0)
    assert snapshot["device_route_distance_meters"] == pytest.approx(41.521)
    assert snapshot["device_speed_meters_per_second"] == pytest.approx(1.375)
    assert snapshot["modeling_reports"] == 1
    assert snapshot["modeling_decode_errors"] == 0


def test_decode_modeling_report_fails_closed_on_invalid_status() -> None:
    with pytest.raises(ModelingReportDecodeError, match="no scan_status"):
        decode_modeling_report(_uint(2, 1))
    with pytest.raises(ModelingReportDecodeError, match="finite and nonnegative"):
        decode_modeling_report(_bytes(3, _fixed32(1, math.nan)))
    with pytest.raises(ModelingReportDecodeError, match="configured limit"):
        decode_modeling_report(b"x" * 5, max_payload_bytes=4)
    with pytest.raises(ModelingReportDecodeError, match="pgo_progress is duplicated"):
        decode_modeling_report(_uint(2, 1) + _uint(2, 2) + _bytes(3, _uint(3, 2)))
    with pytest.raises(ModelingReportDecodeError, match="scan_status is duplicated"):
        decode_modeling_report(_bytes(3, b"") + _bytes(3, b""))
    with pytest.raises(ModelingReportDecodeError, match="move_speed is duplicated"):
        decode_modeling_report(_bytes(3, _fixed32(2, 1.0) + _fixed32(2, 2.0)))
    with pytest.raises(ModelingReportDecodeError, match="signed range"):
        decode_modeling_report(_uint(2, 0x8000_0000) + _bytes(3, _uint(3, 2)))


def test_modeling_metrics_follow_device_scan_generation_reset() -> None:
    metrics = BridgeMetrics()
    before_reset = _bytes(
        3,
        _fixed32(1, 2.071593) + _fixed32(2, 0.2) + _uint(3, 734),
    )
    after_reset = _bytes(
        3,
        _fixed32(1, 0.0) + _fixed32(2, 0.0) + _uint(3, 2),
    )

    observe_modeling_report(
        SimpleNamespace(topic=MODELING_REPORT_TOPIC, payload=before_reset),
        metrics,
    )
    observe_modeling_report(
        SimpleNamespace(topic=MODELING_REPORT_TOPIC, payload=after_reset),
        metrics,
    )

    snapshot = metrics.snapshot()
    assert snapshot["device_elapsed_seconds"] == 1.0
    assert snapshot["device_route_distance_meters"] == 0.0
    assert snapshot["device_speed_meters_per_second"] == 0.0
    assert snapshot["modeling_reports"] == 2


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")
