"""The processor's history stays the acquisition, in order, when the GUI feeds sweeps back.

To recompute the whole B-scan the GUI pushes its retained sweeps into the processor's
input ring again (AppWindowBscanReplayMixin). The processor used to append them to its
own history as if they were new, so its next replay published every sweep two or three
times and out of order, and the B-scan came out with its columns rotated.

Runs the real `data_processor` on private rings; skipped when it has not been built.
"""

from __future__ import annotations

import json
import os
import subprocess
import tempfile
import time
import unittest
from pathlib import Path

import numpy as np

from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.orchestration.live_processing_config import ProcessingLiveConfig, ProcessingLiveConfigWriter
from python_app.orchestration.shm import ShmRingWriter
from python_app.orchestration.shm.ring_reader import ShmRingReader
from python_app.storage.npz.serialize import PREPROC_MAGIC, serialize_trace_collection

_ROOT = Path(__file__).resolve().parents[2]
_PROCESSOR = _ROOT / "build/bin/data_processor"
_SWEEPS = 5
_TIMEOUT_S = 5.0


def _sweep(collection_id: int, monotonic_ns: int) -> SweepCollection:
    points = 64
    trace = TraceData(
        combo=ComboKey(input=0, output=0),
        frequency_hz=np.linspace(1e9, 6e9, points, dtype=np.float32),
        s11=np.zeros(points, dtype=np.complex64),
        s21=np.full(points, collection_id, dtype=np.complex64),
    )
    return SweepCollection(collection_id=collection_id, monotonic_ns=monotonic_ns, traces=[trace])


@unittest.skipUnless(_PROCESSOR.exists(), "data_processor is not built (run make)")
class FedBackSweepsTest(unittest.TestCase):
    def setUp(self) -> None:
        workdir = Path(tempfile.mkdtemp(prefix="processor_history_"))
        config = json.loads((_ROOT / "run_config.json").read_text(encoding="utf-8"))
        for name, ring in config["rings"].items():
            ring["name"] = f"/test_processor_history_{os.getpid()}_{name}"
        config["run"]["combos"] = [{"input": 0, "output": 0}]
        config["run"]["processing_live_config_path"] = str(workdir / "processing_live.json")
        config["run"]["locator_server"]["enabled"] = False
        config_path = workdir / "run_config.json"
        config_path.write_text(json.dumps(config), encoding="utf-8")
        self._live = ProcessingLiveConfig(processor_mode="bscan")
        self._live_writer = ProcessingLiveConfigWriter(workdir / "processing_live.json")
        self._live_writer.write(self._live)

        rings = config["rings"]
        shm_files = [Path("/dev/shm" + ring["name"]) for ring in rings.values()]
        processor = subprocess.Popen(
            [str(_PROCESSOR), "--config", str(config_path)],
            cwd=_ROOT,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

        def cleanup() -> None:
            processor.terminate()
            processor.wait(_TIMEOUT_S)
            for path in shm_files:
                path.unlink(missing_ok=True)
            for path in workdir.iterdir():
                path.unlink()
            workdir.rmdir()

        self.addCleanup(cleanup)
        # The processor creates both rings at start-up, the results ring last.
        self._results = ShmRingReader(rings["results"]["name"], open_timeout_s=_TIMEOUT_S)
        self.addCleanup(self._results.close)
        ring = rings["preprocessed"]
        self._input = ShmRingWriter(ring["name"], int(ring["capacity"]), int(ring["slot_size_bytes"]))
        self.addCleanup(self._input.close)

    def _push(self, sweeps: list[SweepCollection]) -> None:
        for sweep in sweeps:
            self.assertTrue(self._input.push(serialize_trace_collection(sweep, PREPROC_MAGIC)))

    def _published_ids(self, count: int) -> list[int]:
        """Ids of the next results: `count` of them plus any that follow shortly after."""
        ids: list[int] = []
        deadline = time.monotonic() + _TIMEOUT_S
        while time.monotonic() < deadline:
            collection = self._results.pop_result_collection()
            if collection is None:
                time.sleep(0.01)
                continue
            ids.append(int(collection.collection_id))
            if len(ids) == count:
                deadline = min(deadline, time.monotonic() + 0.3)
        return ids

    def test_fed_back_sweeps_are_processed_but_not_retained(self) -> None:
        start_ns = time.monotonic_ns()
        sweeps = [_sweep(index, start_ns + index * 1_000_000) for index in range(1, _SWEEPS + 1)]
        expected = [sweep.collection_id for sweep in sweeps]
        self._push(sweeps)
        self.assertEqual(self._published_ids(_SWEEPS), expected)

        # The GUI feeds them back: each is processed again under its own id.
        self._push(sweeps)
        self.assertEqual(self._published_ids(_SWEEPS), expected)

        # A settings change makes the processor replay its history: each sweep once, in order.
        self._live.bscan_gain = 2.0
        self._live_writer.write(self._live)
        self.assertEqual(self._published_ids(_SWEEPS), expected)


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