"""Decode RoboSense H32 MSOP packets into Cartesian frames (metres).

Angle / distance conventions follow the H32 Medulla plugins:
azimuth = normalize(-(block_az + horizontal[ch])), altitude = vertical[ch],
distance_mm = raw * distance_unit_mm, then:

    x = d_m * cos(alt) * cos(az)
    y = d_m * cos(alt) * sin(az)
    z = d_m * sin(alt)

When DIFOP is unavailable, vertical angles default to a uniform -16°…+16° fan
and horizontal channel offsets default to 0.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import numpy as np

from .capture_format_v2 import CaptureFile

PACKET_LENGTH = 1248
DATA_START = 42
BLOCKS = 12
BLOCK_LENGTH = 100
CHANNELS = 32
MIN_FRAME_POINTS_DEFAULT = 100


def default_vertical_deg() -> np.ndarray:
    return -16.0 + np.arange(CHANNELS, dtype=np.float64) * (32.0 / (CHANNELS - 1))


def default_horizontal_deg() -> np.ndarray:
    return np.zeros(CHANNELS, dtype=np.float64)


def read_u16_be(packet: bytes, index: int) -> int:
    return (packet[index] << 8) | packet[index + 1]


def device_timestamp_ms(packet: bytes) -> int:
    seconds = int.from_bytes(packet[20:26], "big")
    microseconds = int.from_bytes(packet[26:30], "big")
    return seconds * 1000 + microseconds // 1000


def distance_unit_mm(packet: bytes, *, auto: bool = True, fallback: float = 2.5) -> float:
    if not auto:
        return float(fallback)
    return 2.5 if packet[17] == 1 else 0.5


def normalize_azimuth_deg(angle: float) -> float:
    while angle > 180.0:
        angle -= 360.0
    while angle < -180.0:
        angle += 360.0
    return angle


@dataclass
class LidarFrameExport:
    t_start_s: float
    t_end_s: float
    points_xyz: np.ndarray  # (N, 3) metres


def decode_packet_points(
    packet: bytes,
    vertical_deg: np.ndarray,
    horizontal_deg: np.ndarray,
    *,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
) -> tuple[list[float], np.ndarray]:
    """Decode one MSOP packet into block azimuths and concatenated XYZ points."""

    if len(packet) != PACKET_LENGTH:
        return [], np.zeros((0, 3), dtype=np.float64)
    unit = distance_unit_mm(packet)
    az_list: list[float] = []
    chunks: list[np.ndarray] = []
    idx = DATA_START
    for _block in range(BLOCKS):
        if idx + BLOCK_LENGTH > PACKET_LENGTH or packet[idx] != 255 or packet[idx + 1] != 238:
            break
        az = read_u16_be(packet, idx + 2) * 0.01
        az_list.append(az)
        pts = _block_points(
            packet,
            idx,
            az,
            unit,
            vertical_deg,
            horizontal_deg,
            min_range_m=min_range_m,
            max_range_m=max_range_m,
        )
        if pts.shape[0]:
            chunks.append(pts)
        idx += BLOCK_LENGTH
    if not chunks:
        return az_list, np.zeros((0, 3), dtype=np.float64)
    return az_list, np.vstack(chunks)


def _block_points(
    packet: bytes,
    block_offset: int,
    az_deg: float,
    unit_mm: float,
    vertical_deg: np.ndarray,
    horizontal_deg: np.ndarray,
    *,
    min_range_m: float,
    max_range_m: float,
) -> np.ndarray:
    xs: list[float] = []
    ys: list[float] = []
    zs: list[float] = []
    idx = block_offset + 4  # after FF EE + azimuth
    for ch in range(CHANNELS):
        raw = read_u16_be(packet, idx)
        idx += 3
        if raw == 0:
            continue
        d_m = (raw * unit_mm) * 0.001
        if d_m < min_range_m or d_m > max_range_m:
            continue
        az_ch = np.deg2rad(normalize_azimuth_deg(-(az_deg + float(horizontal_deg[ch]))))
        alt = np.deg2rad(float(vertical_deg[ch]))
        cos_alt = np.cos(alt)
        xs.append(d_m * cos_alt * np.cos(az_ch))
        ys.append(d_m * cos_alt * np.sin(az_ch))
        zs.append(d_m * np.sin(alt))
    if not xs:
        return np.zeros((0, 3), dtype=np.float64)
    return np.column_stack([xs, ys, zs]).astype(np.float64, copy=False)


def iter_h32_frames_from_packets(
    packets: Iterable[bytes],
    *,
    min_frame_points: int = MIN_FRAME_POINTS_DEFAULT,
    frame_stride: int = 1,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
    max_points_per_frame: int | None = None,
    vertical_deg: np.ndarray | None = None,
    horizontal_deg: np.ndarray | None = None,
) -> list[LidarFrameExport]:
    """Assemble raw MSOP packets into frames using the 270°→90° azimuth wrap."""

    vertical = default_vertical_deg() if vertical_deg is None else np.asarray(vertical_deg, dtype=np.float64)
    horizontal = default_horizontal_deg() if horizontal_deg is None else np.asarray(horizontal_deg, dtype=np.float64)
    if vertical.shape != (CHANNELS,) or horizontal.shape != (CHANNELS,):
        raise ValueError(f"vertical/horizontal must have shape ({CHANNELS},)")

    frames: list[LidarFrameExport] = []
    point_chunks: list[np.ndarray] = []
    t_start: float | None = None
    t_end: float | None = None
    prev_az: float | None = None
    kept = 0
    stride = max(1, int(frame_stride))

    def emit() -> None:
        nonlocal point_chunks, t_start, t_end, kept
        if not point_chunks or t_start is None or t_end is None:
            point_chunks = []
            t_start = t_end = None
            return
        points = np.vstack(point_chunks)
        point_chunks = []
        start_s, end_s = t_start, t_end
        t_start = t_end = None
        if points.shape[0] < min_frame_points:
            return
        if kept % stride != 0:
            kept += 1
            return
        kept += 1
        if max_points_per_frame is not None and points.shape[0] > max_points_per_frame:
            select = np.linspace(0, points.shape[0] - 1, max_points_per_frame, dtype=int)
            points = points[select]
        if end_s <= start_s:
            end_s = start_s + 0.1
        frames.append(LidarFrameExport(t_start_s=start_s, t_end_s=end_s, points_xyz=points))

    for packet in packets:
        if len(packet) != PACKET_LENGTH:
            continue
        packet_t = device_timestamp_ms(packet) * 1e-3
        unit = distance_unit_mm(packet)
        idx = DATA_START
        for _block in range(BLOCKS):
            if idx + BLOCK_LENGTH > PACKET_LENGTH or packet[idx] != 255 or packet[idx + 1] != 238:
                break
            az = read_u16_be(packet, idx + 2) * 0.01
            if prev_az is not None and prev_az > 270.0 and az < 90.0:
                emit()
            prev_az = az
            pts = _block_points(
                packet,
                idx,
                az,
                unit,
                vertical,
                horizontal,
                min_range_m=min_range_m,
                max_range_m=max_range_m,
            )
            if pts.shape[0]:
                if t_start is None:
                    t_start = packet_t
                t_end = packet_t
                point_chunks.append(pts)
            idx += BLOCK_LENGTH

    emit()
    return frames


def iter_h32_frames(
    capture: CaptureFile,
    *,
    min_frame_points: int = MIN_FRAME_POINTS_DEFAULT,
    frame_stride: int = 1,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
    max_points_per_frame: int | None = None,
    vertical_deg: np.ndarray | None = None,
    horizontal_deg: np.ndarray | None = None,
) -> list[LidarFrameExport]:
    """Assemble MSOP packets from a V2 .rscap capture into frames."""

    return iter_h32_frames_from_packets(
        (chunk.raw for chunk in capture.chunks),
        min_frame_points=min_frame_points,
        frame_stride=frame_stride,
        min_range_m=min_range_m,
        max_range_m=max_range_m,
        max_points_per_frame=max_points_per_frame,
        vertical_deg=vertical_deg,
        horizontal_deg=horizontal_deg,
    )