"""JSON export keeps real stage data and never guesses across repeated sweep IDs."""

from __future__ import annotations

import unittest

import numpy as np

from python_app.models.dataset_model import ComboKey, ResultCollection, SweepCollection, TraceData
from python_app.storage.npz.vna_history_json import build_vna_history_payload


def sweep(cid: int, ns: int, value: float, window: tuple[int, int] = (0, 0)) -> SweepCollection:
    trace = TraceData(
        ComboKey(0, 0), np.array([1e8, 2e8]), np.full(2, value, complex),
        capture_start_ns=window[0], capture_end_ns=window[1],
    )
    return SweepCollection(cid, ns, [trace], capture_start_ns=window[0], capture_end_ns=window[1])


def payload(raw: list[SweepCollection], pre: list[SweepCollection], **kwargs):
    return build_vna_history_payload(raw, pre, [], input_index=0, output_index=0, **kwargs)


class VnaHistoryJsonTest(unittest.TestCase):
    def test_missing_stage_is_empty_instead_of_fabricated(self) -> None:
        for stage in ("raw", "preprocessed"):
            with self.subTest(stage=stage):
                raw, pre = ([sweep(1, 100, 7)], []) if stage == "raw" else ([], [sweep(1, 110, 9)])
                history = payload(raw, pre)["sweep_history"]
                self.assertEqual(len(history), 1)
                self.assertEqual(history[0]["sweep_points"], [[7., 0.], [7., 0.]] if raw else [])
                self.assertEqual(history[0]["calibrated_points"], [[9., 0.], [9., 0.]] if pre else [])

    def test_dropped_stages_and_reused_ids_preserve_every_trace_once(self) -> None:
        raw = [sweep(1, 100, 1, (80, 90)), sweep(1, 200, 2, (180, 190)), sweep(2, 300, 3, (280, 290))]
        pre = [sweep(1, 210, 20, (180, 190)), sweep(3, 410, 40, (380, 390))]
        for stage in ("raw", "preprocessed"):
            with self.subTest(primary=stage):
                history = payload(raw, pre, primary_stage=stage)["sweep_history"]
                self.assertEqual(len(history), 4)
                self.assertEqual([row["sweep_points"][0][0] if row["sweep_points"] else None for row in history], [1, 2, 3, None])
                self.assertEqual([row["calibrated_points"][0][0] if row["calibrated_points"] else None for row in history], [None, 20, None, 40])

    def test_conflicting_windows_are_never_paired(self) -> None:
        history = payload([sweep(1, 100, 1, (80, 90))], [sweep(1, 110, 10, (81, 91))])["sweep_history"]
        self.assertEqual(len(history), 2)
        self.assertEqual(history[0]["calibrated_points"], [])
        self.assertEqual(history[1]["sweep_points"], [])

    def test_export_timestamp_uses_selected_trace_in_a_switched_collection(self) -> None:
        pre = sweep(1, 200, 10, (100, 150))
        pre.capture_start_ns, pre.capture_end_ns = 1, 190
        history = payload([], [pre])["sweep_history"]
        self.assertEqual(history[0]["timestamp"], 150 / 1e9)

    def test_join_uses_collection_window_and_capture_order_not_trace_or_stage_time(self) -> None:
        raw = sweep(1, 900, 1, (10, 20))
        pre = sweep(1, 1000, 10, (10, 20))
        pre.traces[0].capture_start_ns = 11
        pre.traces[0].capture_end_ns = 19
        fresh = sweep(1, 100, 2, (30, 40))
        history = payload([fresh, raw], [pre])["sweep_history"]
        self.assertEqual(len(history), 2)
        self.assertEqual(history[0]["sweep_points"], [[1., 0.], [1., 0.]])
        self.assertEqual(history[0]["calibrated_points"], [[10., 0.], [10., 0.]])
        self.assertEqual(history[0]["timestamp"], 19 / 1e9)
        self.assertEqual(history[1]["calibrated_points"], [])

    def test_reordered_results_do_not_produce_an_alignment_warning(self) -> None:
        pre = [sweep(1, 110, 1, (10, 20)), sweep(2, 210, 2, (30, 40))]
        result = [ResultCollection(2, 999, capture_start_ns=30, capture_end_ns=40),
                  ResultCollection(1, 888, capture_start_ns=10, capture_end_ns=20)]
        exported = build_vna_history_payload([], pre, result, input_index=0, output_index=0)
        self.assertNotIn("alignment_warning", exported)

    def test_selected_combo_with_many_distinct_ids_has_no_lost_or_duplicated_records(self) -> None:
        raw = [sweep(cid, cid * 100, cid, (cid * 100 - 20, cid * 100 - 10)) for cid in range(1, 2001)]
        pre = [sweep(cid, cid * 100 + 10, cid * 2, (cid * 100 - 20, cid * 100 - 10)) for cid in range(1, 2001)]
        history = payload(raw, pre)["sweep_history"]
        self.assertEqual(len(history), 2000)
        self.assertEqual(history[-1]["sweep_points"][0][0], 2000)
        self.assertEqual(history[-1]["calibrated_points"][0][0], 4000)


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