"""Decode Wheeltec N300 FDILink IMU frames from a V2 .rscap capture."""

from __future__ import annotations

import struct
from dataclasses import dataclass

import numpy as np

from .capture_format_v2 import CaptureFile, RawChunk, iter_contiguous_segments


@dataclass(frozen=True)
class ImuSample:
    t_s: float
    gyro_rad_s: tuple[float, float, float]
    accel_m_s2: tuple[float, float, float]
    host_receive_utc_ticks: int
    device_timestamp_us: int


def crc8_fdilink(data: bytes) -> int:
    crc = 0
    for value in data:
        crc ^= value
        for _ in range(8):
            crc = ((crc << 1) ^ 0x07) & 0xFF if crc & 0x80 else (crc << 1) & 0xFF
    return crc


def crc16_fdilink(data: bytes) -> int:
    crc = 0
    for value in data:
        crc ^= value << 8
        for _ in range(8):
            crc = ((crc << 1) ^ 0x1021) & 0xFFFF if crc & 0x8000 else (crc << 1) & 0xFFFF
    return crc


def _host_ticks_for_span(chunks: list[RawChunk], start: int, end: int) -> int:
    stream_offset = 0
    last = chunks[0]
    for chunk in chunks:
        next_offset = stream_offset + len(chunk.raw)
        if start < next_offset and end > stream_offset:
            last = chunk
        stream_offset = next_offset
    return last.receive_utc_ticks


def iter_n300_imu_samples(capture: CaptureFile) -> list[ImuSample]:
    """Return CRC-valid MSG_IMU (0x40) samples sorted by device timestamp."""

    samples: list[ImuSample] = []
    expected_lengths = {0x40: 56, 0x41: 48}
    for _segment_id, chunks in iter_contiguous_segments(capture.chunks):
        stream = b"".join(chunk.raw for chunk in chunks)
        cursor = 0
        while cursor < len(stream):
            start = stream.find(b"\xFC", cursor)
            if start < 0:
                break
            if start + 8 > len(stream):
                break
            payload_length = stream[start + 2]
            end = start + payload_length + 8
            if end > len(stream):
                if stream.find(b"\xFC", start + 1) < 0:
                    break
                cursor = start + 1
                continue
            frame = stream[start:end]
            if frame[-1] != 0xFD:
                cursor = start + 1
                continue
            packet_id = frame[1]
            payload = frame[7:-1]
            header_ok = crc8_fdilink(frame[:4]) == frame[4]
            payload_ok = crc16_fdilink(payload) == int.from_bytes(frame[5:7], "big")
            expected = expected_lengths.get(packet_id)
            length_ok = expected is None or len(payload) == expected
            if not (header_ok and payload_ok and length_ok):
                cursor = start + 1
                continue
            if packet_id == 0x40:
                gyro = struct.unpack_from("<3f", payload, 0)
                accel = struct.unpack_from("<3f", payload, 12)
                device_us = struct.unpack_from("<q", payload, 48)[0]
                samples.append(
                    ImuSample(
                        t_s=float(device_us) * 1e-6,
                        gyro_rad_s=gyro,
                        accel_m_s2=accel,
                        host_receive_utc_ticks=_host_ticks_for_span(chunks, start, end),
                        device_timestamp_us=int(device_us),
                    )
                )
            cursor = end
    samples.sort(key=lambda sample: (sample.t_s, sample.device_timestamp_us))
    return samples


def samples_to_arrays(samples: list[ImuSample]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    if not samples:
        return (
            np.zeros(0, dtype=np.float64),
            np.zeros((0, 3), dtype=np.float64),
            np.zeros((0, 3), dtype=np.float64),
        )
    t = np.asarray([sample.t_s for sample in samples], dtype=np.float64)
    gyro = np.asarray([sample.gyro_rad_s for sample in samples], dtype=np.float64)
    accel = np.asarray([sample.accel_m_s2 for sample in samples], dtype=np.float64)
    return t, gyro, accel