"""Binary decoders for raw/preprocessed/result payload collections."""

from __future__ import annotations

import numpy as np

from python_app.models.dataset_model import (
    ComboKey,
    ResultBlock,
    ResultCollection,
    ResultPayload,
    SweepCollection,
    TraceData,
)
from python_app.orchestration.shm.binary_cursor import ByteCursor

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


def decode_trace_collection(payload: bytes, expected_magic: int) -> SweepCollection:
    """Decode one raw/preprocessed collection from binary payload."""
    cursor = ByteCursor(payload)
    magic = cursor.read_u32()
    if magic != expected_magic:
        raise ValueError("Unexpected trace collection magic")

    collection_id = cursor.read_u64()
    monotonic_ns = cursor.read_u64()
    trace_count = cursor.read_u32()

    traces: list[TraceData] = []
    for _ in range(trace_count):
        input_pos = cursor.read_u32()
        output_pos = cursor.read_u32()
        point_count = cursor.read_u32()

        freq_bytes = point_count * 4
        freq = np.frombuffer(cursor.read_bytes(freq_bytes), dtype="<f4").astype(np.float32, copy=False)

        interleaved_bytes = point_count * 8
        interleaved = np.frombuffer(cursor.read_bytes(interleaved_bytes), dtype="<f4")
        s21 = (interleaved[0::2] + 1j * interleaved[1::2]).astype(np.complex64, copy=False)

        traces.append(
            TraceData(
                combo=ComboKey(input_pos=input_pos, output_pos=output_pos),
                frequency_hz=freq,
                s21=s21,
            )
        )

    return SweepCollection(collection_id=collection_id, monotonic_ns=monotonic_ns, traces=traces)


def decode_result_collection(payload: bytes) -> ResultCollection:
    """Decode one processed result collection from binary payload."""
    cursor = ByteCursor(payload)
    magic = cursor.read_u32()
    if magic != RESULT_MAGIC:
        raise ValueError("Unexpected result collection magic")

    collection_id = cursor.read_u64()
    monotonic_ns = cursor.read_u64()
    block_count = cursor.read_u32()

    blocks: list[ResultBlock] = []
    for _ in range(block_count):
        input_pos = cursor.read_u32()
        output_pos = cursor.read_u32()
        payload_count = cursor.read_u32()

        payloads: list[ResultPayload] = []
        for _ in range(payload_count):
            kind = cursor.read_u8()
            name_size = cursor.read_u16()
            name = cursor.read_bytes(name_size).decode("utf-8")

            if kind == 1:
                point_count = cursor.read_u32()
                freq = np.frombuffer(cursor.read_bytes(point_count * 4), dtype="<f4").astype(np.float32, copy=False)
                interleaved = np.frombuffer(cursor.read_bytes(point_count * 8), dtype="<f4")
                trace = (interleaved[0::2] + 1j * interleaved[1::2]).astype(np.complex64, copy=False)
                payloads.append(
                    ResultPayload(
                        processing_name=name,
                        kind=kind,
                        frequency_hz=freq,
                        trace=trace,
                    )
                )
            elif kind == 2:
                scalar_value = cursor.read_f32()
                payloads.append(
                    ResultPayload(
                        processing_name=name,
                        kind=kind,
                        frequency_hz=np.array([], dtype=np.float32),
                        trace=np.array([], dtype=np.complex64),
                        scalar_value=scalar_value,
                    )
                )
            else:
                raise ValueError(f"Unsupported result payload kind: {kind}")

        blocks.append(
            ResultBlock(
                combo=ComboKey(input_pos=input_pos, output_pos=output_pos),
                payloads=payloads,
            )
        )

    return ResultCollection(collection_id=collection_id, monotonic_ns=monotonic_ns, blocks=blocks)