"""Tests for the Kamil ADC TTY wire-protocol parser."""

from __future__ import annotations

import struct
import unittest

import numpy as np

from python_app.hardware_full.kamil_adc.protocol import (
    MAIN_MARKER,
    REFERENCE_MARKER,
    KamilAdcStreamParser,
)


def _boundary() -> bytes:
    return struct.pack("<HHHH", MAIN_MARKER, 0xFFFF, 0xFFFF, 0xFFFF)


def _main(step: int, real: int, imag: int) -> bytes:
    return struct.pack("<HHhh", MAIN_MARKER, step, real, imag)


def _reference(step: int, real: int, imag: int) -> bytes:
    return struct.pack("<HHhh", REFERENCE_MARKER, step, real, imag)


class KamilAdcStreamParserTest(unittest.TestCase):
    def test_pairs_main_and_reference_by_step(self) -> None:
        parser = KamilAdcStreamParser()
        stream = (
            _boundary()
            + _main(1, 10, -1) + _reference(1, 100, 5)
            + _main(2, 20, -2) + _reference(2, 200, 6)
            + _boundary()
        )
        sweeps = parser.feed(stream)
        self.assertEqual(len(sweeps), 1)
        sweep = sweeps[0]
        self.assertEqual(sweep.steps.tolist(), [1, 2])
        self.assertEqual(sweep.main.tolist(), [complex(10, -1), complex(20, -2)])
        self.assertEqual(sweep.reference.tolist(), [complex(100, 5), complex(200, 6)])

    def test_keeps_only_steps_present_in_both_channels(self) -> None:
        parser = KamilAdcStreamParser()
        stream = (
            _boundary()
            + _main(1, 10, 0)                     # main only -> dropped
            + _main(2, 20, 0) + _reference(2, 200, 0)   # both -> kept
            + _reference(3, 300, 0)               # reference only -> dropped
            + _boundary()
        )
        (sweep,) = parser.feed(stream)
        self.assertEqual(sweep.steps.tolist(), [2])
        self.assertEqual(sweep.main.tolist(), [complex(20, 0)])
        self.assertEqual(sweep.reference.tolist(), [complex(200, 0)])

    def test_orders_by_ascending_step_regardless_of_arrival(self) -> None:
        parser = KamilAdcStreamParser()
        stream = (
            _boundary()
            + _main(3, 3, 0) + _reference(3, 30, 0)
            + _main(1, 1, 0) + _reference(1, 10, 0)
            + _main(2, 2, 0) + _reference(2, 20, 0)
            + _boundary()
        )
        (sweep,) = parser.feed(stream)
        self.assertEqual(sweep.steps.tolist(), [1, 2, 3])
        self.assertEqual(sweep.main.real.tolist(), [1, 2, 3])

    def test_discards_preroll_before_first_boundary(self) -> None:
        parser = KamilAdcStreamParser()
        # Garbage + a partial point before the first real boundary must be skipped.
        stream = (
            _main(7, 7, 7)            # pre-roll point (no preceding boundary) -> ignored
            + _boundary()
            + _main(1, 11, 0) + _reference(1, 1, 0)
            + _boundary()
        )
        (sweep,) = parser.feed(stream)
        self.assertEqual(sweep.steps.tolist(), [1])

    def test_handles_chunk_splits_across_frames(self) -> None:
        parser = KamilAdcStreamParser()
        stream = (
            _boundary()
            + _main(1, 10, 0) + _reference(1, 100, 0)
            + _main(2, 20, 0) + _reference(2, 200, 0)
            + _boundary()
        )
        sweeps: list = []
        # Feed one byte at a time to exercise reassembly across feed() calls.
        for byte in stream:
            sweeps.extend(parser.feed(bytes([byte])))
        self.assertEqual(len(sweeps), 1)
        self.assertEqual(sweeps[0].steps.tolist(), [1, 2])

    def test_multiple_sweeps_in_one_feed(self) -> None:
        parser = KamilAdcStreamParser()
        stream = (
            _boundary()
            + _main(1, 1, 0) + _reference(1, 10, 0)
            + _boundary()
            + _main(1, 2, 0) + _reference(1, 20, 0)
            + _boundary()
        )
        sweeps = parser.feed(stream)
        self.assertEqual(len(sweeps), 2)
        self.assertEqual(sweeps[0].main.real.tolist(), [1])
        self.assertEqual(sweeps[1].main.real.tolist(), [2])

    def test_empty_sweep_between_boundaries_is_skipped(self) -> None:
        parser = KamilAdcStreamParser()
        stream = _boundary() + _boundary() + _main(1, 5, 0) + _reference(1, 1, 0) + _boundary()
        sweeps = parser.feed(stream)
        self.assertEqual(len(sweeps), 1)
        self.assertEqual(sweeps[0].steps.tolist(), [1])

    def test_corrupt_marker_raises(self) -> None:
        parser = KamilAdcStreamParser()
        corrupt = struct.pack("<HHhh", 0x001A, 1, 5, 5)  # unknown marker
        with self.assertRaisesRegex(ValueError, "protocol violation"):
            parser.feed(_boundary() + _main(1, 1, 0) + corrupt + _boundary())

    def test_partial_trailing_frame_is_buffered(self) -> None:
        parser = KamilAdcStreamParser()
        self.assertEqual(parser.feed(_boundary() + _main(1, 1, 0)[:5]), [])
        # Supply the rest of the frame plus its reference and the closing boundary.
        rest = _main(1, 1, 0)[5:]
        (sweep,) = parser.feed(rest + _reference(1, 9, 0) + _boundary())
        self.assertEqual(sweep.steps.tolist(), [1])

    def test_reset_clears_state(self) -> None:
        parser = KamilAdcStreamParser()
        parser.feed(_boundary() + _main(1, 1, 0))
        parser.reset()
        # After reset we must re-align on a fresh boundary before collecting.
        sweeps = parser.feed(_main(9, 9, 0) + _boundary() + _main(1, 2, 0) + _reference(1, 2, 0) + _boundary())
        self.assertEqual(len(sweeps), 1)
        self.assertEqual(sweeps[0].main.real.tolist(), [2])

    def test_dtypes(self) -> None:
        parser = KamilAdcStreamParser()
        (sweep,) = parser.feed(_boundary() + _main(1, 1, 2) + _reference(1, 3, 4) + _boundary())
        self.assertEqual(sweep.main.dtype, np.complex64)
        self.assertEqual(sweep.reference.dtype, np.complex64)
        self.assertEqual(sweep.steps.dtype, np.int32)


if __name__ == "__main__":
    unittest.main()