"""Re-feeding retained sweeps to the processor to recompute the whole B-scan.

The processor keeps only ~50 preprocessed sweeps of its own, so a live settings edit
used to refresh just the newest 50 columns. The GUI holds up to 1000 of them and hands
them back through the processor's input ring, which only works if the bytes the GUI
writes are exactly the ones the C++ side expects — that wire-format contract is what
the first test pins, without needing the pipeline running.

The guard tests pin the other half: the GUI must never write into that ring while
`data_preprocessor` owns it.
"""

from __future__ import annotations

from contextlib import suppress
import os
from pathlib import Path
import unittest

import numpy as np

os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
from PyQt6.QtCore import QSignalBlocker  # noqa: E402
from PyQt6.QtWidgets import QApplication  # noqa: E402

from python_app.gui.app_window import AppWindow  # noqa: E402
from python_app.models.dataset_model import (  # noqa: E402
    ComboKey,
    ResultBlock,
    ResultCollection,
    ResultPayload,
    SweepCollection,
    TraceData,
)
from python_app.orchestration.shm import ShmRingWriter  # noqa: E402
from python_app.orchestration.shm.ring_reader import ShmRingReader  # noqa: E402
from python_app.storage.npz.serialize import PREPROC_MAGIC, serialize_trace_collection  # noqa: E402

# Mirrors kPreprocessedCollectionMagic in common_cpp/ipc/src/shared_types.cpp. If this
# ever drifts, the processor silently rejects everything the GUI re-feeds.
_CPP_PREPROCESSED_MAGIC = 0x32525050

_app: QApplication | None = None
_window: AppWindow | None = None


def setUpModule() -> None:
    global _app, _window
    _app = QApplication.instance() or QApplication([])
    _window = AppWindow(Path("."))


def tearDownModule() -> None:
    if _window is not None:
        _window.close()


def _sweep(collection_id: int, *, combos: int = 3, points: int = 64) -> SweepCollection:
    traces = [
        TraceData(
            combo=ComboKey(input=index, output=index + 1),
            frequency_hz=np.linspace(1e9, 8e9, points, dtype=np.float32),
            s11=(np.arange(points) + index).astype(np.complex64) + 0.5j,
            s21=(np.arange(points) * 2 + index).astype(np.complex64) - 0.25j,
        )
        for index in range(combos)
    ]
    return SweepCollection(
        collection_id=collection_id,
        monotonic_ns=collection_id * 1_000_000 + 7,
        traces=traces,
        capture_start_ns=collection_id * 10,
        capture_end_ns=collection_id * 10 + 5,
    )


def _bscan_result(collection_id: int, *, points: int = 32) -> ResultCollection:
    payload = ResultPayload(
        processing_name="bscan",
        kind=1,
        frequency_hz=np.linspace(0.0, 1.0, points, dtype=np.float32),
        trace=(np.arange(points) + collection_id).astype(np.complex64),
    )
    return ResultCollection(
        collection_id=collection_id,
        monotonic_ns=collection_id,
        blocks=[ResultBlock(combo=ComboKey(input=0, output=0), payloads=[payload])],
    )


class PreprocessedWireFormatTest(unittest.TestCase):
    """The bytes the GUI re-feeds must be the ones the C++ processor decodes."""

    def setUp(self) -> None:
        self.ring_name = f"/radar_test_{self._testMethodName}"
        with suppress(OSError):
            (Path("/dev/shm") / self.ring_name[1:]).unlink()
        self.writer = ShmRingWriter(self.ring_name, 8, 1 << 20)
        self.reader = ShmRingReader(self.ring_name)
        self.addCleanup(self._cleanup)

    def _cleanup(self) -> None:
        with suppress(Exception):
            self.reader.close()
        with suppress(Exception):
            self.writer.close()
        with suppress(OSError):
            (Path("/dev/shm") / self.ring_name[1:]).unlink()

    def _assert_same(self, original: SweepCollection, decoded: SweepCollection) -> None:
        self.assertEqual(decoded.collection_id, original.collection_id)
        self.assertEqual(decoded.monotonic_ns, original.monotonic_ns)
        self.assertEqual(decoded.capture_start_ns, original.capture_start_ns)
        self.assertEqual(decoded.capture_end_ns, original.capture_end_ns)
        self.assertEqual(len(decoded.traces), len(original.traces))
        for left, right in zip(original.traces, decoded.traces, strict=True):
            self.assertEqual((left.combo.input, left.combo.output), (right.combo.input, right.combo.output))
            np.testing.assert_array_equal(left.frequency_hz, right.frequency_hz)
            np.testing.assert_array_equal(left.s11, right.s11)
            np.testing.assert_array_equal(left.s21, right.s21)

    def test_magic_matches_the_cpp_decoder(self) -> None:
        payload = serialize_trace_collection(_sweep(1), PREPROC_MAGIC)
        self.assertEqual(int.from_bytes(payload[:4], "little"), _CPP_PREPROCESSED_MAGIC)

    def test_round_trip_preserves_every_field(self) -> None:
        original = _sweep(42)
        self.assertTrue(self.writer.push(serialize_trace_collection(original, PREPROC_MAGIC)))

        decoded = self.reader.pop_preprocessed_collection()
        self.assertIsNotNone(decoded)
        assert decoded is not None
        self._assert_same(original, decoded)

    def test_batch_keeps_order(self) -> None:
        originals = [_sweep(cid) for cid in range(100, 105)]
        for collection in originals:
            self.assertTrue(self.writer.push(serialize_trace_collection(collection, PREPROC_MAGIC)))

        decoded = []
        while (collection := self.reader.pop_preprocessed_collection()) is not None:
            decoded.append(collection)

        self.assertEqual([c.collection_id for c in decoded], [c.collection_id for c in originals])
        for left, right in zip(originals, decoded, strict=True):
            self._assert_same(left, right)

    def test_overflow_drops_oldest(self) -> None:
        # Why the replay pushes in chunks instead of one burst: an undrained ring
        # silently overwrites, which would punch holes into the very image we are
        # trying to make coherent.
        capacity = 8
        for cid in range(200, 200 + capacity + 4):
            self.writer.push(serialize_trace_collection(_sweep(cid), PREPROC_MAGIC))

        survived = []
        while (collection := self.reader.pop_preprocessed_collection()) is not None:
            survived.append(collection.collection_id)

        self.assertEqual(len(survived), capacity)
        self.assertEqual(survived[-1], 200 + capacity + 3)


class _IdleResultReader:
    """Stands in for a connected results reader that simply has nothing to hand out."""

    def pop_result_collection(self):
        return None

    def close(self) -> None:
        return None


class ReplayGuardTest(unittest.TestCase):
    """The GUI must refuse to write into a ring `data_preprocessor` still owns."""

    def setUp(self) -> None:
        self.w = _window
        # The 50 ms ring poll and the debounce would both run against this half-faked
        # window while the test drives it by hand; park them for the duration.
        self.w._timer.stop()
        self.w._bscan_reprocess_timer.stop()
        self.addCleanup(self.w._timer.start)
        self.addCleanup(self.w._bscan_reprocess_timer.stop)

        self._original_is_running = self.w._supervisor.is_running
        self._original_is_processor_running = self.w._supervisor.is_processor_running
        self._original_mode = self.w._processing_mode.currentText()
        # Changing the mode fires the live-settings handler, which would kick off the
        # very replay these tests are inspecting.
        with QSignalBlocker(self.w._processing_mode):
            self.w._processing_mode.setCurrentText("bscan")
        self.w._supervisor.is_running = lambda: False
        self.w._supervisor.is_processor_running = lambda: True
        self.w._result_reader = _IdleResultReader()

    def tearDown(self) -> None:
        self.w._supervisor.is_running = self._original_is_running
        self.w._supervisor.is_processor_running = self._original_is_processor_running
        with QSignalBlocker(self.w._processing_mode):
            self.w._processing_mode.setCurrentText(self._original_mode)
        self.w._result_reader = None

    def test_refuses_while_acquisition_runs(self) -> None:
        self.w._supervisor.is_running = lambda: True
        self.assertFalse(self.w._can_reprocess_history())

    def test_refuses_when_processor_is_down(self) -> None:
        self.w._supervisor.is_processor_running = lambda: False
        self.assertFalse(self.w._can_reprocess_history())

    def test_refuses_outside_bscan_mode(self) -> None:
        self.w._processing_mode.setCurrentText("pass_through")
        self.assertFalse(self.w._can_reprocess_history())

    def test_refuses_without_a_results_reader(self) -> None:
        self.w._result_reader = None
        self.assertFalse(self.w._can_reprocess_history())

    def test_allows_when_stopped_with_a_live_processor(self) -> None:
        self.assertTrue(self.w._can_reprocess_history())

    def test_replayed_results_are_rendered_instead_of_runtime_history(self) -> None:
        """Regression: 300 re-sent, 300 recovered, only 145 drawn.

        `_pre_history` and `_result_history` come from two rings that drop
        independently, so the re-sent sweeps only partly overlap the recorded results.
        The non-overlapping ones get appended to the deque under old ids, leaving its
        newest `window` entries a mix of freshly and stale-processed frames. Rendering
        from the replayed set instead sidesteps the whole problem.
        """
        window = 30
        replayed = [_bscan_result(cid) for cid in range(9000, 9000 + window)]

        self.w._result_history.clear()
        # Stand-in for a runtime history whose ids barely overlap the replayed ones.
        self.w._result_history.extend(_bscan_result(cid) for cid in range(1, 200))
        self.addCleanup(self.w._result_history.clear)

        self.w._bscan_replay_results = replayed
        self.addCleanup(self.w._discard_bscan_replay_results)
        self.w._bscan_history_window.setValue(window)
        self.w._bscan_render_signature = None
        self.w._sync_bscan_history_from_results()

        self.assertEqual(len(self.w._bscan_history_by_combo[(0, 0)]), window)

    def test_discarding_replay_falls_back_to_runtime_history(self) -> None:
        self.w._bscan_replay_results = [_bscan_result(1)]
        self.w._discard_bscan_replay_results()
        self.assertEqual(self.w._bscan_replay_results, [])
        self.assertIsNone(self.w._bscan_render_signature)

    def test_replay_refuses_to_re_enter_itself(self) -> None:
        # Pumping the event loop mid-replay can fire the debounce again; a nested burst
        # would share the ring and corrupt the result accounting.
        self.w._pre_history.clear()
        self.w._pre_history.extend(_sweep(cid) for cid in range(1, 6))
        self.addCleanup(self.w._pre_history.clear)

        self.w._bscan_replay_active = True
        self.addCleanup(setattr, self.w, "_bscan_replay_active", False)

        self.assertFalse(self.w._reprocess_history_through_processor())

    def test_replay_refuses_without_an_active_run_config(self) -> None:
        # A GUI restarted against a still-running processor does not know the ring
        # geometry, and ShmRingWriter would recreate the segment underneath it.
        self.w._pre_history.clear()
        self.w._pre_history.extend(_sweep(cid) for cid in range(1, 6))
        self.addCleanup(self.w._pre_history.clear)

        previous = getattr(self.w, "_active_run_config", None)
        self.w._active_run_config = None
        self.addCleanup(setattr, self.w, "_active_run_config", previous)

        self.assertIsNone(self.w._ensure_replay_ring_writer())
        self.assertFalse(self.w._reprocess_history_through_processor())


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