"""Binary payload serialization for sweep/result collections."""

from __future__ import annotations

import struct

import numpy as np

from python_app.models.dataset_model import ResultCollection, SweepCollection

RAW_MAGIC = 0x31574152
PREPROC_MAGIC = 0x31525050
RESULT_MAGIC = 0x314C5352


def serialize_trace_collection(collection: SweepCollection, magic: int) -> bytes:
    """Serialize one raw/preprocessed trace collection into ring-compatible binary format."""
    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)
        if freq.size != s21.size:
            raise ValueError("Trace frequency and S21 sizes must match")

        buffer.extend(struct.pack("<III", trace.combo.input_pos, trace.combo.output_pos, int(freq.size)))
        buffer.extend(freq.astype("<f4", copy=False).tobytes())

        interleaved = np.empty(freq.size * 2, dtype="<f4")
        interleaved[0::2] = s21.real.astype("<f4", copy=False)
        interleaved[1::2] = s21.imag.astype("<f4", copy=False)
        buffer.extend(interleaved.tobytes())

    return bytes(buffer)


def serialize_result_collection(collection: ResultCollection) -> bytes:
    """Serialize one processed collection with result blocks/payloads."""
    buffer = bytearray()
    buffer.extend(
        struct.pack("<IQQI", RESULT_MAGIC, collection.collection_id, collection.monotonic_ns, len(collection.blocks))
    )

    for block in collection.blocks:
        buffer.extend(struct.pack("<II", block.combo.input_pos, block.combo.output_pos))
        buffer.extend(struct.pack("<I", len(block.payloads)))

        for payload in block.payloads:
            name_bytes = payload.processing_name.encode("utf-8")
            if len(name_bytes) > 0xFFFF:
                raise ValueError("processing_name is too long")

            buffer.extend(struct.pack("<BH", payload.kind, len(name_bytes)))
            buffer.extend(name_bytes)

            if payload.kind == 1:
                freq = np.asarray(payload.frequency_hz, dtype=np.float32)
                trace = np.asarray(payload.trace, dtype=np.complex64)
                if freq.size != trace.size:
                    raise ValueError("Result trace frequency and values sizes must match")

                buffer.extend(struct.pack("<I", int(freq.size)))
                buffer.extend(freq.astype("<f4", copy=False).tobytes())

                interleaved = np.empty(freq.size * 2, dtype="<f4")
                interleaved[0::2] = trace.real.astype("<f4", copy=False)
                interleaved[1::2] = trace.imag.astype("<f4", copy=False)
                buffer.extend(interleaved.tobytes())
            elif payload.kind == 2:
                buffer.extend(struct.pack("<f", float(payload.scalar_value)))
            else:
                raise ValueError(f"Unsupported payload kind: {payload.kind}")

    return bytes(buffer)