"""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 = 0x32574152
PREPROC_MAGIC = 0x32525050
RESULT_MAGIC = 0x324C5352  # RSL2: adds processing_duration_ns after monotonic_ns


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
        s11_interleaved = np.frombuffer(cursor.read_bytes(interleaved_bytes), dtype="<f4")
        s21_interleaved = np.frombuffer(cursor.read_bytes(interleaved_bytes), dtype="<f4")
        s11 = (s11_interleaved[0::2] + 1j * s11_interleaved[1::2]).astype(np.complex64, copy=False)
        s21 = (s21_interleaved[0::2] + 1j * s21_interleaved[1::2]).astype(np.complex64, copy=False)

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

    # Optional trailer, written after the trace blocks by newer producers: the
    # collection capture window, then a per-trace window table. Both stages are
    # optional so payloads from an older producer still decode (the timestamps
    # simply stay zero).
    capture_start_ns = 0
    capture_end_ns = 0
    if cursor.remaining_bytes() != 0:
        if cursor.remaining_bytes() < 16:
            raise ValueError("Truncated capture window in trace collection")
        capture_start_ns = cursor.read_u64()
        capture_end_ns = cursor.read_u64()

    if cursor.remaining_bytes() != 0:
        trace_time_count = cursor.read_u32()
        if trace_time_count != len(traces):
            raise ValueError("Per-trace capture window count does not match trace count")
        for trace in traces:
            trace.capture_start_ns = cursor.read_u64()
            trace.capture_end_ns = cursor.read_u64()

    if cursor.remaining_bytes() != 0:
        raise ValueError("Unexpected trailing bytes in trace collection")

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


def decode_result_collection(payload: bytes) -> ResultCollection:
    """Decode one processed result collection from binary payload."""

    def read_payload(cursor: ByteCursor) -> ResultPayload:
        """Decode one result payload from stream."""
        kind = cursor.read_u8()
        name_size = cursor.read_u16()
        name = cursor.read_str(name_size)

        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)
            return ResultPayload(
                processing_name=name,
                kind=kind,
                frequency_hz=freq,
                trace=trace,
            )
        if kind == 2:
            return ResultPayload(
                processing_name=name,
                kind=kind,
                scalar_value=cursor.read_f32(),
            )
        if kind == 3:
            x_count = cursor.read_u32()
            y_count = cursor.read_u32()
            image_x_axis = np.frombuffer(cursor.read_bytes(x_count * 4), dtype="<f4").astype(np.float32, copy=False)
            image_y_axis = np.frombuffer(cursor.read_bytes(y_count * 4), dtype="<f4").astype(np.float32, copy=False)
            value_count = x_count * y_count
            image_values = np.frombuffer(cursor.read_bytes(value_count * 4), dtype="<f4").astype(np.float32, copy=False)
            image = image_values.reshape((y_count, x_count))
            return ResultPayload(
                processing_name=name,
                kind=kind,
                image_x_axis=image_x_axis,
                image_y_axis=image_y_axis,
                image=image,
            )
        if kind == 4:
            table_columns = cursor.read_u32()
            table_rows = cursor.read_u32()
            if table_rows and not table_columns:
                raise ValueError("Result table with rows must have at least one column")
            value_count = table_columns * table_rows
            table_values = np.frombuffer(cursor.read_bytes(value_count * 4), dtype="<f4").astype(np.float32, copy=False)
            table = table_values.reshape((table_rows, table_columns))
            return ResultPayload(
                processing_name=name,
                kind=kind,
                table=table,
            )
        raise ValueError(f"Unsupported result payload kind: {kind}")

    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()
    processing_duration_ns = cursor.read_u64()
    collection_payload_count = cursor.read_u32()
    block_count = cursor.read_u32()

    collection_payloads: list[ResultPayload] = []
    for _ in range(collection_payload_count):
        collection_payloads.append(read_payload(cursor))

    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):
            payloads.append(read_payload(cursor))

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

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