"""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 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_last_cycle_s_parameters,
)
from python_app.hardware_full.librevna_multi_device_driver.exceptions import TransientCollectionError
from python_app.hardware_full.librevna_multi_device_driver.models import normalize_master_stimulus_ports
from python_app.hardware_full.librevna_multi_device_driver.protocol import parse_vna_datapoint_payload
from python_app.hardware_full.multi_device_service import (
    _MAX_TRANSIENT_COLLECTION_RETRIES,
    MultiDeviceLibreVnaService,
)


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 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)
        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 LastCycleSParameterTest(unittest.TestCase):
    """Reference/measurement division, gap zero-filling, and the over-loss guard.

    A handful of dropped USB frames (points that never arrived, left as ``NaN``)
    are emitted as zero so the cycle is kept instead of stalling on recovery, while
    delivered points keep their true normalized value. Losing more than
    ``_MAX_ZERO_FILLED_POINTS_FRACTION`` of a trace is treated as a genuine fault
    and rejected as transient. The tiny-trace cases below cross that fraction with a
    single missing point, so they still raise.
    """

    def test_divides_measurement_by_reference(self) -> None:
        reference = {1: np.array([[2 + 0j, 0 + 2j, 4 + 0j]], dtype=complex)}
        measurement = {"S31": np.array([[4 + 0j, 0 + 4j, 2 + 0j]], dtype=complex)}
        result = calculate_last_cycle_s_parameters(measurement, reference)
        np.testing.assert_allclose(result["s31"], np.array([2 + 0j, 2 + 0j, 0.5 + 0j]))

    def test_single_dropped_measurement_is_zero_filled(self) -> None:
        # One missing point in a realistically sized trace (< the loss fraction) is
        # emitted as zero; every delivered point keeps its real value, no extras.
        point_count = 200
        reference = {1: np.full((1, point_count), 2 + 0j, dtype=complex)}
        measurement = {"S31": np.full((1, point_count), 4 + 0j, dtype=complex)}
        measurement["S31"][0, 100] = np.nan + 1j * np.nan
        result = calculate_last_cycle_s_parameters(measurement, reference)
        self.assertEqual(result["s31"][100], 0)
        np.testing.assert_allclose(np.delete(result["s31"], 100), 2 + 0j)

    def test_single_dropped_reference_is_zero_filled(self) -> None:
        # A missing master reference cannot normalize that point for any trace, so
        # the point is zeroed; the rest of the trace is untouched.
        point_count = 200
        reference = {1: np.full((1, point_count), 2 + 0j, dtype=complex)}
        reference[1][0, 50] = np.nan + 1j * np.nan
        measurement = {"S31": np.full((1, point_count), 4 + 0j, dtype=complex)}
        result = calculate_last_cycle_s_parameters(measurement, reference)
        self.assertEqual(result["s31"][50], 0)
        np.testing.assert_allclose(np.delete(result["s31"], 50), 2 + 0j)

    def test_single_near_zero_reference_is_zero_filled(self) -> None:
        # A ~zero incident reference would explode the division into noise; one such
        # point is zeroed rather than aborting the whole cycle, and stays finite.
        point_count = 200
        reference = {1: np.full((1, point_count), 2 + 0j, dtype=complex)}
        reference[1][0, 75] = 0 + 0j
        measurement = {"S31": np.full((1, point_count), 4 + 0j, dtype=complex)}
        result = calculate_last_cycle_s_parameters(measurement, reference)
        self.assertEqual(result["s31"][75], 0)
        self.assertTrue(np.isfinite(result["s31"]).all())

    def test_excessive_loss_is_transient(self) -> None:
        # Losing more than the tolerated fraction is a real fault, not a stray
        # dropped frame: reject as transient so the caller recovers properly.
        point_count = 200
        reference = {1: np.full((1, point_count), 2 + 0j, dtype=complex)}
        measurement = {"S31": np.full((1, point_count), 4 + 0j, dtype=complex)}
        measurement["S31"][0, :40] = np.nan + 1j * np.nan  # 20% > 5% tolerance
        with self.assertRaises(TransientCollectionError):
            calculate_last_cycle_s_parameters(measurement, reference)

    def test_nan_reference_is_transient(self) -> None:
        reference = {1: np.array([[2 + 0j, np.nan + 1j * np.nan, 4 + 0j]], dtype=complex)}
        measurement = {"S31": np.array([[4 + 0j, 4 + 0j, 2 + 0j]], dtype=complex)}
        with self.assertRaises(TransientCollectionError):
            calculate_last_cycle_s_parameters(measurement, reference)

    def test_nan_measurement_is_transient(self) -> None:
        reference = {1: np.array([[2 + 0j, 2 + 0j, 4 + 0j]], dtype=complex)}
        measurement = {"S31": np.array([[4 + 0j, np.nan + 1j * np.nan, 2 + 0j]], dtype=complex)}
        with self.assertRaises(TransientCollectionError):
            calculate_last_cycle_s_parameters(measurement, reference)

    def test_near_zero_reference_is_transient(self) -> None:
        reference = {1: np.array([[2 + 0j, 0 + 0j, 4 + 0j]], dtype=complex)}
        measurement = {"S31": np.array([[4 + 0j, 4 + 0j, 2 + 0j]], dtype=complex)}
        with self.assertRaises(TransientCollectionError):
            calculate_last_cycle_s_parameters(measurement, reference)


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."""

    def _acquire_native_collection(self, collection_id: int, capture_start_ns: int):  # type: ignore[override]
        self.collection_calls += 1
        return self.behavior(collection_id, capture_start_ns)

    def recover(self, **_kwargs) -> None:  # type: ignore[override]
        self.reopen_calls += 1


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

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

    def test_succeeds_after_transient_failures_without_reopen(self) -> None:
        def flaky(_collection_id: int, _capture_start_ns: int) -> str:
            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, 0)
        self.assertEqual(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, _capture_start_ns: int) -> str:
            raise TransientCollectionError("persistent drop")

        service = self._stub(always_transient)
        with self.assertRaises(TransientCollectionError):
            service._acquire_native_collection_with_transient_retries(1, 0)
        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, _capture_start_ns: int) -> str:
            raise RuntimeError("USB transport failed")

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


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