"""Unit tests for the LibreVNA multi-device acquisition pipeline.

These cover the host-side logic that does not need real hardware: frame resync
(F2), datapoint parse hardening (N2), the cross-device cycle-0 alignment check
(N1), and the transient-vs-fatal error classification driving the retry tier
(F1/F3).
"""

from __future__ import annotations

import struct
import threading
import unittest
from unittest.mock import Mock, patch
from collections.abc import Callable

from python_app.models.dataset_model import SweepCollection

from python_app.hardware_full.librevna_driver.enums import PacketType as NativePacketType
from python_app.hardware_full.librevna_driver.models import Packet
from python_app.hardware_full.librevna_driver.protocol.frame import FrameScanner, encode_frame
import numpy as np

from python_app.hardware_full.librevna_multi_device_driver.cycle_collection import (
    _CrossDeviceCycleSynchronizer,
    _expected_sweep_cycle_seconds,
    _is_transient_collection_error,
    calculate_s_parameters,
    collect_complete_running_sweep_cycle,
)
from python_app.hardware_full.librevna_multi_device_driver.exceptions import TransientCollectionError
from python_app.hardware_full.librevna_multi_device_driver.models import (
    SweepConfiguration,
    build_output_channels,
    normalize_master_stimulus_ports,
)
from python_app.hardware_full.librevna_multi_device_driver.protocol import (
    DISABLED_PORT_STAGE,
    build_sweep_settings_payload,
    parse_vna_datapoint_payload,
)
from python_app.hardware_full.multi_device_service import (
    _MAX_TRANSIENT_COLLECTION_RETRIES,
    MultiDeviceLibreVnaService,
)
from python_app.models.run_config_schema import RadarSweepModel


def _control_frame(packet_type: NativePacketType, payload: bytes) -> bytes:
    """Encode a CRC-checked control frame (e.g. ACK)."""
    return encode_frame(Packet(packet_type, payload))


def _datapoint_frame(frequency_hz: int, point_index: int, values: list[tuple[int, complex]]) -> bytes:
    """Encode a CRC-less VNADatapoint frame from (description_mask, value) pairs."""
    masks = bytes(mask for mask, _ in values)
    reals = struct.pack(f"<{len(values)}f", *[value.real for _, value in values])
    imags = struct.pack(f"<{len(values)}f", *[value.imag for _, value in values])
    payload = struct.pack("<QhH", int(frequency_hz), 0, int(point_index)) + reals + imags + masks
    return encode_frame(Packet(NativePacketType.VNA_DATAPOINT, payload))


class FrameScannerResyncTest(unittest.TestCase):
    """F2: the scanner must resync past corruption instead of raising/dying."""

    def test_decodes_clean_back_to_back_frames(self) -> None:
        scanner = FrameScanner()
        stream = _control_frame(NativePacketType.ACK, b"") + _control_frame(NativePacketType.SET_IDLE, b"")
        packets = scanner.feed(stream)
        self.assertEqual([p.type for p in packets], [NativePacketType.ACK, NativePacketType.SET_IDLE])
        self.assertEqual(scanner.discarded_byte_count, 0)

    def test_resyncs_past_garbage_between_valid_frames(self) -> None:
        scanner = FrameScanner()
        good_a = _control_frame(NativePacketType.ACK, b"\x01\x02\x03")
        good_b = _control_frame(NativePacketType.SET_IDLE, b"")
        # Garbage that contains stray 0x5A bytes (false headers) with junk lengths.
        garbage = bytes([0x5A, 0x10, 0x00, 0xFF, 0x5A, 0x5A, 0xAB, 0xCD])
        packets = scanner.feed(good_a + garbage + good_b)
        self.assertEqual([p.type for p in packets], [NativePacketType.ACK, NativePacketType.SET_IDLE])
        self.assertGreater(scanner.discarded_byte_count, 0)

    def test_resyncs_on_corrupted_crc(self) -> None:
        scanner = FrameScanner()
        good = _control_frame(NativePacketType.ACK, b"")
        corrupt = bytearray(_control_frame(NativePacketType.SET_IDLE, b"\x07\x07"))
        corrupt[-1] ^= 0xFF  # flip a CRC byte -> CRCError on decode
        recovered = _control_frame(NativePacketType.ACK, b"\x09")
        packets = scanner.feed(good + bytes(corrupt) + recovered)
        # The corrupt frame is dropped; the surrounding valid frames survive.
        self.assertEqual([p.type for p in packets], [NativePacketType.ACK, NativePacketType.ACK])
        self.assertGreater(scanner.discarded_byte_count, 0)

    def test_never_raises_on_pure_garbage(self) -> None:
        scanner = FrameScanner()
        # A blob full of 0x5A false headers must never raise and never hang.
        blob = bytes((0x5A, 0x00, 0xFF, 0x5A, 0x08, 0x00, 0x63) * 200)
        packets = scanner.feed(blob)
        self.assertEqual(packets, [])

    def test_reassembles_frame_split_across_chunks(self) -> None:
        scanner = FrameScanner()
        frame = _control_frame(NativePacketType.ACK, b"\xaa\xbb\xcc\xdd")
        self.assertEqual(scanner.feed(frame[:3]), [])
        packets = scanner.feed(frame[3:])
        self.assertEqual([p.type for p in packets], [NativePacketType.ACK])
        self.assertEqual(scanner.discarded_byte_count, 0)

    def test_datapoint_frame_round_trips_without_crc(self) -> None:
        scanner = FrameScanner()
        frame = _datapoint_frame(3_000_000_000, 0, [(0x00, 1 + 2j), (0x10, 3 - 4j)])
        packets = scanner.feed(frame)
        self.assertEqual([p.type for p in packets], [NativePacketType.VNA_DATAPOINT])

    def test_corrupted_datapoint_resyncs_to_next_good_frame(self) -> None:
        scanner = FrameScanner()
        # A datapoint frame whose trailing CRC field is non-zero is rejected
        # (VNADatapoint must carry zero CRC); the scanner must resync to the ACK.
        bad = bytearray(_datapoint_frame(3_000_000_000, 5, [(0x00, 1 + 1j)]))
        bad[-1] = 0x01  # non-zero CRC trailer -> CRCError
        good = _control_frame(NativePacketType.ACK, b"")
        packets = scanner.feed(bytes(bad) + good)
        self.assertEqual([p.type for p in packets], [NativePacketType.ACK])
        self.assertGreater(scanner.discarded_byte_count, 0)


class CaptureWindowTest(unittest.TestCase):
    def test_native_capture_window_excludes_discarded_tail_and_queue_drain(self) -> None:
        for first_point in (1, 3):
            with self.subTest(first_point=first_point):
                clock = threading.local()
                configuration = SweepConfiguration(1_000_000, 2_000_000, 5, 1000, -10.0)

                def connection(serial):
                    sequence = iter(enumerate([*range(first_point, 5), *range(5)]))

                    def receive(*, timeout_seconds):
                        offset, point = next(sequence)
                        clock.ns = 1_000_000_000 + offset * 20_000_000
                        values = [(0x13, 2 + 0j), (0x01, 1 + 0j), (0x02, 1 + 0j),
                                  (0x33, 2 + 0j), (0x21, 1 + 0j), (0x22, 1 + 0j)]
                        return 27, _datapoint_frame(1_000_000 + point * 250_000, point, values)[4:-4]

                    return Mock(serial_number=serial, receive_packet=receive)

                master, slave1, slave2 = [connection(serial) for serial in ("M", "A", "B")]
                service = MultiDeviceLibreVnaService("M", ["A", "B"])
                service.configure(RadarSweepModel(start_hz=1_000_000, stop_hz=2_000_000, points=5, if_bandwidth_hz=1000))
                service._controller = Mock(collect_running_sweep_cycle=lambda *_args, **_kwargs: collect_complete_running_sweep_cycle(
                    master_device_connection=master, slave_device_connections=[slave1, slave2],
                    active_sweep_configuration=configuration, master_stimulus_ports=(1, 2),
                ))
                with patch("python_app.hardware_full.multi_device_service.time.monotonic_ns", side_effect=lambda: getattr(clock, "ns", 1)):
                    collection = service.acquire_collection(7)
                start = 1_000_000_000 + (5 - first_point) * 20_000_000
                self.assertEqual((collection.capture_start_ns, collection.capture_end_ns), (start, start + 80_000_000))
                self.assertEqual(len(collection.traces), 8)
                for trace in collection.traces:
                    self.assertEqual((trace.capture_start_ns, trace.capture_end_ns), (start, start + 80_000_000))
                    np.testing.assert_array_equal(trace.s21, np.full(5, 0.5))


class CycleClosureTest(unittest.TestCase):
    def test_next_sweep_without_its_point_zero_still_closes_the_cycle(self) -> None:
        # A point of this sweep and point 0 of the next are lost: the next sweep's point 1
        # must close the cycle instead of overwriting this sweep's point 1.
        points = 40
        configuration = SweepConfiguration(1_000_000, 2_000_000, points, 1000, -10.0)
        frames = [*((point, 1 + 0j) for point in range(points) if point != 3),
                  *((point, 7 + 0j) for point in range(1, points))]

        def connection(serial: str) -> Mock:
            sequence = iter(frames)

            def receive(*, timeout_seconds):
                point, measured = next(sequence)
                values = [(0x13, 2 + 0j), (0x01, measured), (0x02, measured),
                          (0x33, 2 + 0j), (0x21, measured), (0x22, measured)]
                return 27, _datapoint_frame(1_000_000 + point * 25_000, point, values)[4:-4]

            return Mock(serial_number=serial, receive_packet=receive)

        result = collect_complete_running_sweep_cycle(
            master_device_connection=connection("M"), slave_device_connections=[connection("A"), connection("B")],
            active_sweep_configuration=configuration, master_stimulus_ports=(1, 2),
        )
        expected = np.full(points, .5)
        expected[3] = 0  # the lost point is zero-filled
        for trace in result.s_parameters.values():
            np.testing.assert_array_equal(trace, expected)


class CrossDeviceCycleSynchronizerTest(unittest.TestCase):
    """N1: cycle-0 anchors are accepted only when all devices saw the same wrap."""

    @staticmethod
    def _run(device_count: int, anchor_times: list[float], tolerance: float, timeout: float = 2.0):
        sync = _CrossDeviceCycleSynchronizer(
            device_count=device_count,
            alignment_tolerance_seconds=tolerance,
            rendezvous_timeout_seconds=timeout,
        )
        results: dict[int, bool] = {}
        lock = threading.Lock()

        def worker(index: int) -> None:
            verdict = sync.confirm_aligned_cycle_zero(anchor_times[index])
            with lock:
                results[index] = verdict

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(len(anchor_times))]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join(timeout=5.0)
        return results

    def test_aligned_anchors_accept(self) -> None:
        # Three devices saw point 0 within ~0.4 ms (host jitter) -> aligned.
        results = self._run(3, [100.0000, 100.0002, 100.0004], tolerance=2e-3)
        self.assertEqual(results, {0: True, 1: True, 2: True})

    def test_off_by_one_sweep_is_rejected(self) -> None:
        # One device anchored ~8 ms (one sweep period) later -> misaligned.
        results = self._run(3, [100.000, 100.001, 100.008], tolerance=2e-3)
        self.assertEqual(results, {0: False, 1: False, 2: False})

    def test_single_device_trivially_aligns(self) -> None:
        results = self._run(1, [100.0], tolerance=2e-3)
        self.assertEqual(results, {0: True})

    def test_missing_peer_times_out_to_not_aligned(self) -> None:
        # Synchronizer expects 3 devices but only 2 arrive: the barrier breaks
        # after the rendezvous timeout and the waiters report not-aligned.
        results = self._run(3, [100.0, 100.0], tolerance=2e-3, timeout=0.3)
        self.assertEqual(results, {0: False, 1: False})


class MasterStimulusPortValidationTest(unittest.TestCase):
    """Q1: the single shared validator enforces the ports 1/2, unique rule."""

    def test_accepts_valid(self) -> None:
        self.assertEqual(normalize_master_stimulus_ports([1, 2]), (1, 2))
        self.assertEqual(normalize_master_stimulus_ports((2,)), (2,))

    def test_rejects_invalid(self) -> None:
        for bad in ([], [3], [1, 2, 1], [1, 1], [0]):
            with self.assertRaises(ValueError):
                normalize_master_stimulus_ports(bad)


class DatapointParseHardeningTest(unittest.TestCase):
    """N2: reject misframed/misaligned datapoint payloads instead of trusting them."""

    @staticmethod
    def _payload(point_index: int, values: list[tuple[int, complex]]) -> bytes:
        masks = bytes(mask for mask, _ in values)
        reals = struct.pack(f"<{len(values)}f", *[v.real for _, v in values])
        imags = struct.pack(f"<{len(values)}f", *[v.imag for _, v in values])
        return struct.pack("<QhH", 3_000_000_000, 0, point_index) + reals + imags + masks

    def test_valid_payload_parses(self) -> None:
        parsed = parse_vna_datapoint_payload(self._payload(7, [(0x00, 1 + 2j), (0x10, 3 + 4j)]))
        self.assertIsNotNone(parsed)
        assert parsed is not None
        self.assertEqual(parsed.point_index, 7)
        self.assertEqual(parsed.frequency_hz, 3_000_000_000)
        self.assertEqual(set(parsed.receiver_values_by_description_mask), {0x00, 0x10})

    def test_too_short_payload_rejected(self) -> None:
        self.assertIsNone(parse_vna_datapoint_payload(b"\x00" * 8))

    def test_misaligned_body_rejected(self) -> None:
        # 12-byte header + a body that is not a multiple of 9 bytes -> misframe.
        for extra in (1, 5, 8, 10):
            self.assertIsNone(
                parse_vna_datapoint_payload(struct.pack("<QhH", 3_000_000_000, 0, 0) + b"\x00" * extra),
                msg=f"body of {extra} bytes must be rejected",
            )

    def test_zero_values_rejected(self) -> None:
        self.assertIsNone(parse_vna_datapoint_payload(struct.pack("<QhH", 3_000_000_000, 0, 0)))


class CycleAlignmentToleranceCalibrationTest(unittest.TestCase):
    """N1: verify the alignment tolerance is well-separated from jitter and one sweep."""

    def test_matches_measured_751pt_2stage_50khz(self) -> None:
        # Experimental anchor: 751 points x 2 stages at 50 kHz IF -> a real ~150 ms
        # sweep. The conservative estimate is 30 ms and the tolerance is half of it.
        cycle_seconds = _expected_sweep_cycle_seconds(751, 2, 50_000)
        self.assertAlmostEqual(cycle_seconds, 0.03004, places=4)
        tolerance = 0.5 * cycle_seconds
        self.assertAlmostEqual(tolerance, 0.01502, places=4)
        # A true off-by-one is one real sweep (~150 ms): flagged with large margin.
        real_sweep_seconds = 0.150
        self.assertGreater(real_sweep_seconds, 5.0 * tolerance)
        # Typical host jitter (~1 ms) clears the tolerance with room to spare.
        typical_jitter_seconds = 0.001
        self.assertLess(typical_jitter_seconds, 0.25 * tolerance)

    def test_scales_with_stage_count(self) -> None:
        single = _expected_sweep_cycle_seconds(751, 1, 50_000)
        dual = _expected_sweep_cycle_seconds(751, 2, 50_000)
        self.assertAlmostEqual(dual, 2.0 * single, places=6)


class SParameterTest(unittest.TestCase):
    """Normalize delivered samples, zero isolated losses, reject excessive loss."""

    def test_divides_measurement_by_reference(self) -> None:
        result = calculate_s_parameters(
            {(0, 0): np.array([4, 4j, 2])}, {0: np.array([2, 2j, 4])})
        np.testing.assert_allclose(result[(0, 0)], [2, 2, .5])

    def test_lost_measurement_or_unusable_reference(self) -> None:
        for source, value in (("measurement", np.nan), ("reference", np.nan), ("reference", 0)):
            for lost in (1, 40):
                with self.subTest(source=source, value=value, lost=lost):
                    reference, measurement = np.full(200, 2 + 0j), np.full(200, 4 + 0j)
                    (measurement if source == "measurement" else reference)[:lost] = value
                    if lost == 40:
                        with self.assertRaises(TransientCollectionError):
                            calculate_s_parameters({(0, 0): measurement}, {0: reference})
                    else:
                        result = calculate_s_parameters({(0, 0): measurement}, {0: reference})[(0, 0)]
                        np.testing.assert_array_equal(result, np.r_[0, np.full(199, 2)])


class TransientErrorClassificationTest(unittest.TestCase):
    """F1/F3: only device-healthy failures are transient (cheap retry)."""

    def test_transient_types(self) -> None:
        self.assertTrue(_is_transient_collection_error(TransientCollectionError("x")))
        self.assertTrue(_is_transient_collection_error(TimeoutError("no progress")))

    def test_fatal_types(self) -> None:
        self.assertFalse(_is_transient_collection_error(RuntimeError("USB transport failed")))
        self.assertFalse(_is_transient_collection_error(ValueError("bad config")))

    def test_transient_is_runtimeerror_subclass(self) -> None:
        # So existing `except Exception`/`except RuntimeError` recovery still catches it.
        self.assertIsInstance(TransientCollectionError("x"), RuntimeError)


class _StubService(MultiDeviceLibreVnaService):
    """Non-slots subclass so a test can inject collection behaviour + count reopens."""

    behavior: Callable[[int], SweepCollection]
    collection_calls: int = 0
    reopen_calls: int = 0

    def _acquire_native_collection(self, collection_id: int) -> SweepCollection:
        self.collection_calls += 1
        return self.behavior(collection_id)

    def recover(self, **_kwargs) -> None:
        self.reopen_calls += 1


class TransientRetryTierTest(unittest.TestCase):
    """F1/F3: transient failures retry in place (no USB reopen); fatals escalate."""

    @staticmethod
    def _stub(behavior: Callable[[int], SweepCollection]) -> _StubService:
        service = _StubService(master_serial="M", slave_serials=["A", "B"])
        service.behavior = behavior
        service.collection_calls = 0
        service.reopen_calls = 0
        return service

    def test_succeeds_after_transient_failures_without_reopen(self) -> None:
        collection = SweepCollection(1, 1, [])

        def flaky(_collection_id: int) -> SweepCollection:
            if service.collection_calls <= _MAX_TRANSIENT_COLLECTION_RETRIES:
                raise TransientCollectionError("dropped datapoint")
            return collection

        service = self._stub(flaky)
        result = service._acquire_native_collection_with_transient_retries(1)
        self.assertIs(result, collection)
        self.assertEqual(service.collection_calls, _MAX_TRANSIENT_COLLECTION_RETRIES + 1)
        self.assertEqual(service.reopen_calls, 0)  # never reopened USB for transient errors

    def test_exhausted_transient_retries_raise_transient(self) -> None:
        def always_transient(_collection_id: int) -> SweepCollection:
            raise TransientCollectionError("persistent drop")

        service = self._stub(always_transient)
        with self.assertRaises(TransientCollectionError):
            service._acquire_native_collection_with_transient_retries(1)
        self.assertEqual(service.collection_calls, _MAX_TRANSIENT_COLLECTION_RETRIES + 1)

    def test_fatal_error_is_not_retried_in_place(self) -> None:
        def fatal(_collection_id: int) -> SweepCollection:
            raise RuntimeError("USB transport failed")

        service = self._stub(fatal)
        with self.assertRaises(RuntimeError):
            service._acquire_native_collection_with_transient_retries(1)
        self.assertEqual(service.collection_calls, 1)  # fatal escalates immediately, no in-place retry


class PerPointStimulusSwitchingTest(unittest.TestCase):
    """Per-point stimulus switching: one output per (master port, switch position).

    Without hardware these pin the contract the modified master bitstream and the
    trigger-loop switch controller have to agree with.
    """

    def test_disabled_switching_keeps_the_stock_layout(self) -> None:
        channels = build_output_channels((1, 2), 1)
        self.assertEqual(
            [(c.index, c.stimulus_port, c.stage_index) for c in channels],
            [(0, 1, 0), (1, 2, 1)],
        )

    def test_two_positions_give_four_outputs(self) -> None:
        channels = build_output_channels((1, 2), 2)
        # Port-major in BOTH numbers: the output index is the device stage, so outputs
        # come in the order the device measures them — port 1 through its switch
        # positions, then port 2. Verified against a scope on the switch line.
        self.assertEqual(
            [(c.index, c.stimulus_port, c.stage_index) for c in channels],
            [(0, 1, 0), (1, 1, 1), (2, 2, 2), (3, 2, 3)],
        )

    def test_output_index_is_the_device_stage(self) -> None:
        # The whole point of the 2026-09-17 renumbering: an operator hanging antennas
        # sees stage order on a scope, and the output numbering must not permute it.
        for substage_count in (1, 2, 4):
            channels = build_output_channels((1, 2), substage_count)
            self.assertEqual(
                [c.index for c in channels],
                [c.stage_index for c in channels],
                msg=f"substages={substage_count}",
            )
            self.assertEqual([c.index for c in channels], list(range(len(channels))))

    def test_switch_position_is_the_low_stage_bit(self) -> None:
        # This is what lets the controller just invert its pin on every trigger cycle
        # instead of tracking an absolute stage number.
        channels = build_output_channels((1, 2), 2)
        for channel in channels:
            expected_position = channel.index & 1
            self.assertEqual(channel.stage_index & 1, expected_position)

    def test_stage_budget_is_enforced(self) -> None:
        build_output_channels((1, 2), 4)  # 8 stages exactly fits
        with self.assertRaises(ValueError):
            build_output_channels((1, 2), 8)  # 16 stages does not
        for bad in (0, 3, -1):
            with self.assertRaises(ValueError):
                build_output_channels((1, 2), bad)

    def test_sweep_settings_carry_stages_and_substage_width(self) -> None:
        configuration = SweepConfiguration(
            start_hz=1_000_000, stop_hz=2_000_000, points=101, if_bandwidth=1000, power_dbm=-10.0
        )
        master = build_sweep_settings_payload(
            configuration,
            is_synchronization_master=True,
            synchronization_enabled=True,
            master_stimulus_ports=(1, 2),
            switch_substage_count=2,
        )
        stage_word = struct.unpack_from("<H", master, 25)[0]
        self.assertEqual(stage_word & 0x07, 3)            # 4 stages -> last index 3
        # Port stages are the FIRST stage of each port's block: port 1 owns stages
        # 0..1 and port 2 owns 2..3, so the bitstream matches them on the stage bits
        # above the substage block.
        self.assertEqual((stage_word >> 3) & 0x07, 0)
        self.assertEqual((stage_word >> 6) & 0x07, 2)
        self.assertEqual((stage_word >> 9) & 0x07, 1)     # log2(2 positions) = 1 substage bit

        slave = build_sweep_settings_payload(
            configuration,
            is_synchronization_master=False,
            synchronization_enabled=True,
            master_stimulus_ports=(1, 2),
            switch_substage_count=2,
        )
        slave_stage_word = struct.unpack_from("<H", slave, 25)[0]
        self.assertEqual(slave_stage_word & 0x07, 3)      # same stage count as the master
        # Receivers must never excite. The stock "one past the last stage" sentinel is
        # unusable once the port comes from the upper stage bits, so both ports are
        # parked at the explicit disabled value instead.
        self.assertEqual((slave_stage_word >> 3) & 0x07, DISABLED_PORT_STAGE)
        self.assertEqual((slave_stage_word >> 6) & 0x07, DISABLED_PORT_STAGE)

    def test_disabled_switching_sends_the_stock_stage_word(self) -> None:
        configuration = SweepConfiguration(
            start_hz=1_000_000, stop_hz=2_000_000, points=101, if_bandwidth=1000, power_dbm=-10.0
        )
        payload = build_sweep_settings_payload(
            configuration,
            is_synchronization_master=True,
            synchronization_enabled=True,
            master_stimulus_ports=(1, 2),
        )
        stage_word = struct.unpack_from("<H", payload, 25)[0]
        self.assertEqual(stage_word & 0x07, 1)            # 2 stages
        self.assertEqual((stage_word >> 3) & 0x07, 0)
        self.assertEqual((stage_word >> 6) & 0x07, 1)
        self.assertEqual((stage_word >> 9) & 0x07, 0)     # no substage bits -> stock bitstream

    def test_service_reports_four_outputs_when_switching(self) -> None:
        service = MultiDeviceLibreVnaService(
            master_serial="M", slave_serials=["A", "B"], switch_substage_count=2
        )
        self.assertEqual(service.output_position_count, 4)


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