"""Binary payload serialization for raw/preprocessed collections."""

from __future__ import annotations

import struct
from numbers import Integral

import numpy as np

from python_app.models.dataset_model import SweepCollection


def validate_trace_collection(
    collection: SweepCollection, *, require_finite: bool = False, require_capture_window: bool = False,
) -> None:
    """Reject ambiguous combinations and malformed traces before saving any files."""
    for field in ("collection_id", "monotonic_ns", "capture_start_ns", "capture_end_ns"):
        _validate_u64(getattr(collection, field), field)
    if require_capture_window and not 0 < collection.capture_start_ns <= collection.capture_end_ns:
        raise ValueError("Raw collection requires a positive, ordered capture window")
    if require_finite and not collection.traces:
        raise ValueError("Preprocess set must contain at least one trace")
    seen: set[tuple[int, int]] = set()
    for trace in collection.traces:
        _validate_u64(trace.capture_start_ns, "trace.capture_start_ns")
        _validate_u64(trace.capture_end_ns, "trace.capture_end_ns")
        if any(
            isinstance(value, bool) or not isinstance(value, Integral) or not 0 <= value <= 0xFFFFFFFF
            for value in (trace.combo.input, trace.combo.output)
        ):
            raise ValueError("Trace input/output positions must be unsigned 32-bit integers")
        combo = (int(trace.combo.input), int(trace.combo.output))
        if combo in seen:
            raise ValueError(f"Invalid or duplicate trace combination: {combo}")
        seen.add(combo)
        arrays = (
            np.asarray(trace.frequency_hz, dtype=np.float32),
            np.asarray(trace.s21, dtype=np.complex64),
        )
        if any(array.ndim != 1 for array in arrays) or len({array.shape for array in arrays}) != 1:
            raise ValueError(f"Trace {combo} frequency and S21 must be equally sized 1D arrays")
        if require_finite and (not arrays[0].size or any(not np.isfinite(array).all() for array in arrays)):
            raise ValueError(f"Preprocess trace {combo} must contain finite, nonempty arrays")


def _validate_u64(value: int, field: str) -> None:
    if isinstance(value, bool) or not isinstance(value, Integral) or not 0 <= value < 1 << 64:
        raise ValueError(f"{field} must be an unsigned 64-bit integer")


def serialize_trace_collection(
    collection: SweepCollection, magic: int, *, require_capture_window: bool = False,
) -> bytes:
    """Serialize one raw/preprocessed trace collection into ring-compatible binary format.

    Each trace stores its frequency array followed by its complex S21 array, using
    12 bytes per point. The trailer contains the collection capture window and a
    per-trace window table (one ``(start_ns, end_ns)`` pair per trace, in trace
    order). See :func:`python_app.orchestration.shm.decoder.decode_trace_collection`
    and ``read_trace_collection`` in ``common_cpp/ipc/src/shared_types.cpp``.
    """
    validate_trace_collection(collection, require_capture_window=require_capture_window)
    buffer = bytearray()
    buffer.extend(struct.pack("<IQQI", magic, collection.collection_id, collection.monotonic_ns, len(collection.traces)))

    for trace in collection.traces:
        freq = np.asarray(trace.frequency_hz, dtype=np.float32)
        s21 = np.asarray(trace.s21, dtype=np.complex64)
        buffer.extend(struct.pack("<III", trace.combo.input, trace.combo.output, int(freq.size)))
        buffer.extend(freq.astype("<f4", copy=False).tobytes())

        buffer.extend(s21.astype("<c8", copy=False).tobytes())

    buffer.extend(
        struct.pack(
            "<QQ",
            int(collection.capture_start_ns),
            int(collection.capture_end_ns),
        )
    )
    buffer.extend(struct.pack("<I", len(collection.traces)))
    for trace in collection.traces:
        buffer.extend(
            struct.pack("<QQ", int(trace.capture_start_ns), int(trace.capture_end_ns))
        )
    return bytes(buffer)