"""Unit tests for HI13 / HI91 IMU decoding."""

from __future__ import annotations

import struct

from tools.rscap_v2.capture_format_v2 import CaptureFile, CaptureHeader, RawChunk
from tools.rscap_v2.hi13_imu import (
    crc16_hi13,
    iter_hi13_imu_samples,
    parse_hi91_frame,
    parse_hi91_sample,
)


def _hi91_frame(
    *,
    device_ms: int = 123456,
    accel_g=(0.0, 0.0, 1.0),
    gyro_dps=(1.0, -2.0, 3.0),
    rpy_deg=(4.0, 5.0, 6.0),
    quaternion_wxyz=(1.0, 0.0, 0.0, 0.0),
) -> bytes:
    payload = bytearray(76)
    payload[0] = 0x91
    struct.pack_into("<H", payload, 1, 0)  # pps
    payload[3] = 25  # temp
    struct.pack_into("<f", payload, 4, 101325.0)
    struct.pack_into("<I", payload, 8, device_ms)
    struct.pack_into("<fff", payload, 12, *accel_g)
    struct.pack_into("<fff", payload, 24, *gyro_dps)
    struct.pack_into("<fff", payload, 48, *rpy_deg)
    struct.pack_into("<ffff", payload, 60, *quaternion_wxyz)
    payload_length = len(payload)
    header = bytearray(6)
    header[0] = 0x5A
    header[1] = 0xA5
    header[2] = payload_length & 0xFF
    header[3] = (payload_length >> 8) & 0xFF
    frame_wo_crc = bytes(header[:4]) + bytes(payload)
    # crc over header[0:4] + payload
    tmp = bytearray(6 + payload_length)
    tmp[0:4] = header[0:4]
    tmp[6:] = payload
    crc = crc16_hi13(tmp, payload_length)
    header[4] = crc & 0xFF
    header[5] = (crc >> 8) & 0xFF
    return bytes(header) + bytes(payload)


def test_parse_hi91_units():
    frame = _hi91_frame(device_ms=5000, accel_g=(0.0, 0.0, 1.0), gyro_dps=(57.2957795, 0.0, 0.0))
    parsed = parse_hi91_frame(frame)
    assert parsed is not None
    gyro, accel, device_ms = parsed
    assert device_ms == 5000
    assert abs(accel[2] - 9.80665) < 1e-4
    assert abs(gyro[0] - 1.0) < 1e-5
    full = parse_hi91_sample(frame, host_receive_utc_ticks=123)
    assert full is not None
    assert full.system_time_ms == 5000
    assert full.host_receive_utc_ticks == 123
    assert full.rpy_deg == (4.0, 5.0, 6.0)
    assert full.quaternion_wxyz == (1.0, 0.0, 0.0, 0.0)


def test_iter_hi13_from_capture():
    frame = _hi91_frame(device_ms=42)
    header = CaptureHeader(
        sensor_kind="hi13r4-imu",
        session_id="t",
        session_start_utc_ticks=0,
        session_start_monotonic_ticks=0,
        monotonic_frequency=10_000_000,
        port="COM1",
        baud=115200,
        file_start_utc_ticks=0,
    )
    chunk = RawChunk(
        sequence=1,
        receive_utc_ticks=100,
        receive_monotonic_ticks=1,
        raw=frame,
        record_file_offset=0,
        raw_file_offset=0,
        record_crc32=0,
        crc_valid=True,
    )
    capture = CaptureFile(path="mem", header=header, chunks=[chunk], footer=None)
    samples = iter_hi13_imu_samples(capture)
    assert len(samples) == 1
    assert samples[0].device_timestamp_us == 42_000
    assert abs(samples[0].t_s - 0.042) < 1e-12