"""Wire protocol for the Kamil ADC collector's TTY stream.

The collector publishes a stream of 8-byte little-endian frames, four 16-bit
words each: ``[marker, step, ch1, ch2]`` (``ch1``/``ch2`` signed). Three frame
kinds appear in ``do8_freq_ref`` mode:

* **Sweep boundary** — ``marker, 0xFFFF, 0xFFFF, 0xFFFF``. Delimits sweeps.
* **Main point** — ``0x000A, step, I, Q``. One complex main sample at ``step``.
* **Reference point** — ``0x00A8, step, I, Q``. The reference sample paired with
  the main sample of the same ``step`` (emitted only where the DI8 loopback
  settled, so reference points are sparser than main points).

:class:`KamilAdcStreamParser` consumes raw bytes incrementally and yields one
:class:`RawSweep` per completed sweep. Main and reference points are aligned by
step index; only steps carrying *both* survive (a step needs its reference for
the frequency axis and its main for the signal). Anything other than the three
known frame kinds is a protocol violation and raises — the reader fails fast and
lets the supervisor relaunch a clean collector rather than silently resyncing.
"""

from __future__ import annotations

from dataclasses import dataclass
import struct

import numpy as np

# Wire-format constants.
FRAME_BYTES = 8
MAIN_MARKER = 0x000A
REFERENCE_MARKER = 0x00A8
# Combo tag — ``0x00C0, input_pos, output_pos, dirty``. Emitted by the switch-aware
# collector right after a sweep boundary to label the upcoming sweep with the RF
# switch combination it was captured under (``dirty != 0`` means the sweep straddled
# a switch transition and must be dropped). Absent in the standalone/calibration
# collector, where every sweep carries no combo (``combo is None``).
COMBO_MARKER = 0x00C0
_BOUNDARY_STEP = 0xFFFF

# marker (u16), step (u16), ch1 (i16), ch2 (i16) — point frames carry signed I/Q.
_FRAME = struct.Struct("<HHhh")
# The last three words of a boundary frame are all 0xFFFF; a real step is
# 1..0xFFFE, so this tail unambiguously marks a sweep boundary regardless of the
# (profile-dependent) start marker.
_BOUNDARY_TAIL = b"\xff\xff\xff\xff\xff\xff"
# Frame-alignment anchor: the phase-profile sweep-boundary frame.
_START_FRAME = struct.pack("<HHHH", MAIN_MARKER, 0xFFFF, 0xFFFF, 0xFFFF)


@dataclass(frozen=True, slots=True)
class RawSweep:
    """One sweep's main and reference samples, aligned and ordered by step.

    ``steps`` is strictly ascending; ``main`` and ``reference`` are the complex
    samples at those steps. All three arrays share the same length.
    """

    steps: np.ndarray
    main: np.ndarray
    reference: np.ndarray
    combo: tuple[int, int] | None = None
    dirty: bool = False

    @property
    def size(self) -> int:
        return int(self.steps.size)


class KamilAdcStreamParser:
    """Incremental, frame-aligning parser turning the TTY byte stream into sweeps.

    Usage: call :meth:`feed` with each chunk of bytes; it returns the list of
    sweeps completed by that chunk (usually zero or one). The parser is stateful
    but holds no I/O and is cheap to unit-test.
    """

    __slots__ = ("_buffer", "_aligned", "_main", "_reference", "_pending_combo", "_pending_dirty")

    def __init__(self) -> None:
        self._buffer = bytearray()
        self._aligned = False
        self._main: dict[int, complex] = {}
        self._reference: dict[int, complex] = {}
        # Combo tag for the sweep currently being accumulated (set by the combo
        # frame right after each boundary; ``None`` in non-switch collector modes).
        self._pending_combo: tuple[int, int] | None = None
        self._pending_dirty = False

    def feed(self, data: bytes) -> list[RawSweep]:
        """Append ``data`` and return any sweeps completed by it."""
        self._buffer.extend(data)
        if not self._aligned and not self._align():
            return []

        sweeps: list[RawSweep] = []
        buffer = self._buffer
        while len(buffer) >= FRAME_BYTES:
            frame = bytes(buffer[:FRAME_BYTES])
            del buffer[:FRAME_BYTES]
            if frame[2:] == _BOUNDARY_TAIL:
                sweep = self._take_sweep()
                if sweep is not None:
                    sweeps.append(sweep)
                continue
            marker, step, real, imag = _FRAME.unpack(frame)
            if marker == MAIN_MARKER:
                self._main[step] = complex(real, imag)
            elif marker == REFERENCE_MARKER:
                self._reference[step] = complex(real, imag)
            elif marker == COMBO_MARKER:
                # step = input_pos, real = output_pos, imag = dirty flag.
                self._pending_combo = (int(step), int(real))
                self._pending_dirty = imag != 0
            else:
                raise ValueError(
                    f"Kamil ADC protocol violation: unexpected frame marker 0x{marker:04x}"
                )
        return sweeps

    def reset(self) -> None:
        """Drop all buffered state (e.g. after the collector is relaunched)."""
        self._buffer.clear()
        self._aligned = False
        self._main.clear()
        self._reference.clear()
        self._pending_combo = None
        self._pending_dirty = False

    def _align(self) -> bool:
        """Discard pre-roll up to and including the first sweep boundary.

        Returns ``True`` once frame-aligned. Keeps a short tail so a boundary
        split across two feeds can still be found on the next chunk.
        """
        index = self._buffer.find(_START_FRAME)
        if index < 0:
            if len(self._buffer) >= FRAME_BYTES:
                del self._buffer[: -(FRAME_BYTES - 1)]
            return False
        del self._buffer[: index + FRAME_BYTES]
        self._main.clear()
        self._reference.clear()
        self._pending_combo = None
        self._pending_dirty = False
        self._aligned = True
        return True

    def _take_sweep(self) -> RawSweep | None:
        """Assemble the buffered points into a sweep and reset for the next one."""
        shared = sorted(self._main.keys() & self._reference.keys())
        main = self._main
        reference = self._reference
        combo = self._pending_combo
        dirty = self._pending_dirty
        self._main = {}
        self._reference = {}
        # The next sweep's combo is set by its own combo frame (right after this
        # boundary); clear so a sweep without one reports combo=None rather than
        # inheriting a stale tag.
        self._pending_combo = None
        self._pending_dirty = False
        if not shared:
            return None
        return RawSweep(
            steps=np.asarray(shared, dtype=np.int32),
            main=np.asarray([main[step] for step in shared], dtype=np.complex64),
            reference=np.asarray([reference[step] for step in shared], dtype=np.complex64),
            combo=combo,
            dirty=dirty,
        )