"""Parser and CSV renderer for the Trikke TRK1 binary telemetry stream."""

from __future__ import annotations

import math
import struct
import zlib
from dataclasses import dataclass

MAGIC = b"TRK1"
VERSION = 1
HEADER_SIZE = 36
SAMPLE_RECORD_SIZE = 20
METADATA_SIZE = 48
MAX_RECORDS = 8

PACKET_TYPE_METADATA = 1
PACKET_TYPE_SAMPLES = 2
PACKET_FLAG_TIMESTAMP_DELTA_SATURATED = 0x01

HEADER = struct.Struct("<4sBBBBBBHIQIII")
METADATA = struct.Struct("<HHHH10f")
SAMPLE = struct.Struct("<IHhhhhhhBB")

CSV_COLUMNS = [
    "sequence",
    "poll_timestamp_us",
    "accel_x_raw",
    "accel_y_raw",
    "accel_z_raw",
    "gyro_x_raw",
    "gyro_y_raw",
    "gyro_z_raw",
    "accel_x_mg",
    "accel_y_mg",
    "accel_z_mg",
    "gyro_x_mdps",
    "gyro_y_mdps",
    "gyro_z_mdps",
    "accel_native_x_raw",
    "accel_native_y_raw",
    "accel_native_z_raw",
    "gyro_native_x_raw",
    "gyro_native_y_raw",
    "gyro_native_z_raw",
    "accel_int_source",
    "gyro_status",
    "loop_overrun_count",
]


@dataclass(frozen=True)
class Metadata:
    sample_rate_hz: int
    accel_range_g: int
    gyro_range_dps: int
    flags: int
    accel_offset_counts: tuple[float, float, float]
    accel_counts_per_g: tuple[float, float, float]
    gyro_bias_counts: tuple[float, float, float]
    gyro_mdps_per_lsb: float


@dataclass(frozen=True)
class Sample:
    sequence: int
    timestamp_us: int
    accel: tuple[int, int, int]
    gyro: tuple[int, int, int]
    accel_status: int
    gyro_status: int


@dataclass(frozen=True)
class Frame:
    packet_type: int
    flags: int
    packet_sequence: int
    base_timestamp_us: int
    dropped_sample_count: int
    loop_overrun_count: int
    metadata: Metadata | None
    samples: tuple[Sample, ...]
    raw: bytes


@dataclass
class IntegrityTracker:
    packet_gap_count: int = 0
    packet_reset_count: int = 0
    sample_gap_count: int = 0
    sample_reset_count: int = 0
    timing_anomaly_count: int = 0
    timestamp_saturation_frame_count: int = 0
    accel_stale_count: int = 0
    accel_overrun_count: int = 0
    gyro_stale_count: int = 0
    gyro_overrun_count: int = 0
    final_dropped_sample_count: int = 0
    final_loop_overrun_count: int = 0
    _previous_packet_sequence: int | None = None
    _previous_sample_sequence: int | None = None
    _previous_timestamp_us: int | None = None

    @staticmethod
    def _classify_sequence(
        previous: int,
        current: int,
    ) -> tuple[int, int]:
        expected = (previous + 1) & 0xFFFFFFFF
        forward_distance = (current - expected) & 0xFFFFFFFF
        if forward_distance == 0:
            return (0, 0)
        if forward_distance < 0x80000000:
            return (forward_distance, 0)
        return (0, 1)

    def observe(self, frame: Frame) -> None:
        if self._previous_packet_sequence is not None:
            gaps, resets = self._classify_sequence(
                self._previous_packet_sequence, frame.packet_sequence
            )
            self.packet_gap_count += gaps
            self.packet_reset_count += resets
        self._previous_packet_sequence = frame.packet_sequence

        if frame.flags & PACKET_FLAG_TIMESTAMP_DELTA_SATURATED:
            self.timestamp_saturation_frame_count += 1
        self.final_dropped_sample_count = frame.dropped_sample_count
        self.final_loop_overrun_count = frame.loop_overrun_count

        for sample in frame.samples:
            if self._previous_sample_sequence is not None:
                gaps, resets = self._classify_sequence(
                    self._previous_sample_sequence, sample.sequence
                )
                self.sample_gap_count += gaps
                self.sample_reset_count += resets
            if (
                self._previous_timestamp_us is not None
                and sample.timestamp_us - self._previous_timestamp_us != 10_000
            ):
                self.timing_anomaly_count += 1
            if not sample.accel_status & 0x80:
                self.accel_stale_count += 1
            if sample.accel_status & 0x01:
                self.accel_overrun_count += 1
            if not sample.gyro_status & 0x08:
                self.gyro_stale_count += 1
            if sample.gyro_status & 0x80:
                self.gyro_overrun_count += 1
            self._previous_sample_sequence = sample.sequence
            self._previous_timestamp_us = sample.timestamp_us


def _lround(value: float) -> int:
    """Match C lroundf: nearest integer, halfway cases away from zero."""
    return math.floor(value + 0.5) if value >= 0 else math.ceil(value - 0.5)


def sample_to_csv_row(sample: Sample, metadata: Metadata, loop_overruns: int) -> list[int]:
    accel_mg = [
        _lround((sample.accel[i] - metadata.accel_offset_counts[i]) * 1000.0 /
                metadata.accel_counts_per_g[i])
        for i in range(3)
    ]
    gyro_mdps = [
        _lround((sample.gyro[i] - metadata.gyro_bias_counts[i]) *
                metadata.gyro_mdps_per_lsb)
        for i in range(3)
    ]

    # Inverse of enclosure accel(x,y,z)=(native_y,-native_x,native_z).
    accel_native = (-sample.accel[1], sample.accel[0], sample.accel[2])
    gyro_native = sample.gyro
    return [
        sample.sequence,
        sample.timestamp_us,
        *sample.accel,
        *sample.gyro,
        *accel_mg,
        *gyro_mdps,
        *accel_native,
        *gyro_native,
        sample.accel_status,
        sample.gyro_status,
        loop_overruns,
    ]


class StreamParser:
    def __init__(self) -> None:
        self._buffer = bytearray()
        self.synchronized = False
        self.skipped_bytes = 0
        self.header_errors = 0
        self.startup_crc_errors = 0
        self.crc_errors = 0

    @property
    def buffered_bytes(self) -> int:
        """Bytes retained because they do not yet form a complete frame."""
        return len(self._buffer)

    def feed(self, data: bytes) -> list[Frame]:
        self._buffer.extend(data)
        frames: list[Frame] = []
        while True:
            magic_at = self._buffer.find(MAGIC)
            if magic_at < 0:
                keep = min(len(self._buffer), len(MAGIC) - 1)
                self.skipped_bytes += len(self._buffer) - keep
                if keep:
                    del self._buffer[:-keep]
                else:
                    self._buffer.clear()
                break
            if magic_at:
                self.skipped_bytes += magic_at
                del self._buffer[:magic_at]
            if len(self._buffer) < HEADER_SIZE:
                break

            fields = HEADER.unpack_from(self._buffer)
            (
                _magic,
                version,
                packet_type,
                header_size,
                record_size,
                record_count,
                flags,
                payload_size,
                packet_sequence,
                base_timestamp_us,
                dropped_sample_count,
                loop_overrun_count,
                expected_crc,
            ) = fields

            valid_shape = (
                version == VERSION
                and header_size == HEADER_SIZE
                and packet_type in (PACKET_TYPE_METADATA, PACKET_TYPE_SAMPLES)
                and payload_size <= max(METADATA_SIZE, SAMPLE_RECORD_SIZE * MAX_RECORDS)
            )
            if packet_type == PACKET_TYPE_METADATA:
                valid_shape = valid_shape and (
                    record_size == 0 and record_count == 0 and payload_size == METADATA_SIZE
                )
            elif packet_type == PACKET_TYPE_SAMPLES:
                valid_shape = valid_shape and (
                    record_size == SAMPLE_RECORD_SIZE
                    and 1 <= record_count <= MAX_RECORDS
                    and payload_size == record_size * record_count
                )
            if not valid_shape:
                self.header_errors += 1
                self.skipped_bytes += 1
                del self._buffer[0]
                continue

            frame_size = HEADER_SIZE + payload_size
            if len(self._buffer) < frame_size:
                break
            raw = bytes(self._buffer[:frame_size])
            actual_crc = zlib.crc32(raw[4:32])
            actual_crc = zlib.crc32(raw[HEADER_SIZE:], actual_crc)
            if actual_crc != expected_crc:
                if self.synchronized:
                    self.crc_errors += 1
                else:
                    self.startup_crc_errors += 1
                self.skipped_bytes += 1
                del self._buffer[0]
                continue

            metadata = None
            samples: tuple[Sample, ...] = ()
            payload = raw[HEADER_SIZE:]
            if packet_type == PACKET_TYPE_METADATA:
                values = METADATA.unpack(payload)
                metadata = Metadata(
                    sample_rate_hz=values[0],
                    accel_range_g=values[1],
                    gyro_range_dps=values[2],
                    flags=values[3],
                    accel_offset_counts=values[4:7],
                    accel_counts_per_g=values[7:10],
                    gyro_bias_counts=values[10:13],
                    gyro_mdps_per_lsb=values[13],
                )
            else:
                decoded: list[Sample] = []
                timestamp_us = base_timestamp_us
                for index in range(record_count):
                    values = SAMPLE.unpack_from(payload, index * SAMPLE_RECORD_SIZE)
                    if index:
                        timestamp_us += values[1] * 10
                    decoded.append(
                        Sample(
                            sequence=values[0],
                            timestamp_us=timestamp_us,
                            accel=values[2:5],
                            gyro=values[5:8],
                            accel_status=values[8],
                            gyro_status=values[9],
                        )
                    )
                samples = tuple(decoded)

            frames.append(
                Frame(
                    packet_type=packet_type,
                    flags=flags,
                    packet_sequence=packet_sequence,
                    base_timestamp_us=base_timestamp_us,
                    dropped_sample_count=dropped_sample_count,
                    loop_overrun_count=loop_overrun_count,
                    metadata=metadata,
                    samples=samples,
                    raw=raw,
                )
            )
            self.synchronized = True
            del self._buffer[:frame_size]
        return frames