"""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 = 0x32574152
PREPROC_MAGIC = 0x32525050
RESULT_MAGIC = 0x314C5352


def _write_interleaved_complex(buffer: bytearray, values: np.ndarray) -> None:
    """Append complex64 array as interleaved float32 real/imag pairs."""
    interleaved = np.empty(values.size * 2, dtype="<f4")
    interleaved[0::2] = values.real.astype("<f4", copy=False)
    interleaved[1::2] = values.imag.astype("<f4", copy=False)
    buffer.extend(interleaved.tobytes())


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)
        s11 = np.asarray(trace.s11, dtype=np.complex64)
        s21 = np.asarray(trace.s21, dtype=np.complex64)
        if freq.size != s11.size:
            raise ValueError("Trace frequency and S11 sizes must match")
        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())

        _write_interleaved_complex(buffer, s11)
        _write_interleaved_complex(buffer, s21)

    return bytes(buffer)


def serialize_result_collection(collection: ResultCollection) -> bytes:
    """Serialize one processed collection with result blocks/payloads."""

    def serialize_payload(buffer: bytearray, payload) -> None:
        """Append one payload in ring-compatible result format."""
        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())
            return

        if payload.kind == 2:
            buffer.extend(struct.pack("<f", float(payload.scalar_value)))
            return

        if payload.kind == 3:
            image_x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
            image_y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
            image = np.asarray(payload.image, dtype=np.float32)
            if image.ndim != 2:
                raise ValueError("Result image payload must be a 2D matrix")
            if image.shape != (image_y_axis.size, image_x_axis.size):
                raise ValueError("Result image axis sizes must match image matrix shape")
            buffer.extend(struct.pack("<II", int(image_x_axis.size), int(image_y_axis.size)))
            buffer.extend(image_x_axis.astype("<f4", copy=False).tobytes())
            buffer.extend(image_y_axis.astype("<f4", copy=False).tobytes())
            buffer.extend(image.astype("<f4", copy=False).ravel(order="C").tobytes())
            return

        if payload.kind == 4:
            table = np.asarray(payload.table, dtype=np.float32)
            if table.ndim != 2:
                raise ValueError("Result table payload must be a 2D matrix")
            buffer.extend(struct.pack("<II", int(table.shape[1]), int(table.shape[0])))
            buffer.extend(table.astype("<f4", copy=False).ravel(order="C").tobytes())
            return

        raise ValueError(f"Unsupported payload kind: {payload.kind}")

    buffer = bytearray()
    buffer.extend(
        struct.pack(
            "<IQQII",
            RESULT_MAGIC,
            collection.collection_id,
            collection.monotonic_ns,
            len(collection.collection_payloads),
            len(collection.blocks),
        )
    )

    for payload in collection.collection_payloads:
        serialize_payload(buffer, payload)

    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:
            serialize_payload(buffer, payload)

    return bytes(buffer)