"""Neutral pair creation must report incomplete sets and roll back failed writes."""

from __future__ import annotations

from pathlib import Path
import tempfile
import unittest
from unittest import mock

import numpy as np

from python_app.models.run_config_model import RunConfigModel
from python_app.scripts import create_neutral_preprocess_sets as cli
from python_app.storage.npz_store import NpzStore
from python_app.workflows.kamil_adc_neutral_preprocess import build_neutral_s21_sets


class NeutralPreprocessCliTest(unittest.TestCase):
    def setUp(self) -> None:
        self.directory = self.enterContext(tempfile.TemporaryDirectory())
        self.root = Path(self.directory)
        self.store = NpzStore(self.root)
        self.config = RunConfigModel.from_dict({"radar": {
            "model": "librevna", "driver_mode": "mock", "sweep": {"points": 7},
        }})
        calibration, reference = build_neutral_s21_sets(self.config)
        calibration.traces[0].s21[:] = 3 + 1j
        reference.traces[0].s21[:] = 2 - 1j
        self.previous = {"s21_calibration": calibration, "s21_reference": reference}

    def _main(self, *options: str) -> int:
        with (
            mock.patch.object(cli.RunConfigModel, "load_from_path", return_value=self.config),
            mock.patch.object(cli, "radar_key_from_config", return_value="radar"),
            mock.patch.object(cli, "NpzStore", return_value=self.store),
            mock.patch("sys.argv", ["neutral", "--config", "profile.json", "--set-name", "neutral", *options]),
        ):
            return cli.main()

    def _files(self) -> dict[str, bytes]:
        return {str(path.relative_to(self.root)): path.read_bytes() for path in self.root.rglob("*") if path.is_file()}

    def _fail_second_archive(self):
        save_archive = np.savez
        count = 0

        def write(*args, **kwargs):
            nonlocal count
            count += 1
            if count == 2:
                raise OSError("reference write failed")
            return save_archive(*args, **kwargs)

        return mock.patch("python_app.storage.npz.store.np.savez", side_effect=write)

    def test_partial_existing_pair_reports_missing_kind_without_overwriting(self) -> None:
        for present, missing in [("s21_calibration", "s21_reference"), ("s21_reference", "s21_calibration")]:
            with self.subTest(present=present):
                self.store = NpzStore(self.root / present)
                self.store.save_set(present, "radar", "neutral", self.previous[present])
                before = self._files()
                with self.assertRaisesRegex(RuntimeError, f"missing {missing}; use --force to replace both"):
                    self._main()
                self.assertEqual(self._files(), before)

    def test_complete_existing_pair_is_skipped(self) -> None:
        self.store.save_sets("radar", "neutral", self.previous)
        before = self._files()
        with mock.patch.object(self.store, "save_sets", wraps=self.store.save_sets) as save:
            self.assertEqual(self._main(), 0)
        save.assert_not_called()
        self.assertEqual(self._files(), before)

    def test_force_replaces_both_existing_sets_with_neutral_data(self) -> None:
        self.store.save_sets("radar", "neutral", self.previous)
        self.assertEqual(self._main("--force"), 0)
        for kind, value in [("s21_calibration", 1), ("s21_reference", 0)]:
            np.testing.assert_array_equal(self.store.load_set(kind, "radar", "neutral").traces[0].s21, value)
        self.assertEqual(list(self.root.glob(".sets-*")), [])

    def test_force_completes_either_partial_pair(self) -> None:
        for present in self.previous:
            with self.subTest(present=present):
                self.store = NpzStore(self.root / present)
                self.store.save_set(present, "radar", "neutral", self.previous[present])
                self.assertEqual(self._main("--force"), 0)
                for kind, value in [("s21_calibration", 1), ("s21_reference", 0)]:
                    np.testing.assert_array_equal(self.store.load_set(kind, "radar", "neutral").traces[0].s21, value)

    def test_second_archive_failure_restores_previous_pair_byte_for_byte(self) -> None:
        self.store.save_sets("radar", "neutral", self.previous)
        before = self._files()
        with self._fail_second_archive(), self.assertRaisesRegex(OSError, "reference write failed"):
            self._main("--force")
        self.assertEqual(self._files(), before)

    def test_second_archive_failure_removes_new_pair(self) -> None:
        with self._fail_second_archive(), self.assertRaisesRegex(OSError, "reference write failed"):
            self._main()
        self.assertEqual(self._files(), {})

    def test_failed_rollback_keeps_the_previous_archive_for_recovery(self) -> None:
        self.store.save_sets("radar", "neutral", self.previous)
        previous_archive = (self.root / "s21_calibration/radar/neutral.npz").read_bytes()
        replace = Path.replace

        def replace_file(path: Path, destination: Path):
            if path.parent.name.startswith(".sets-") and path.name == "s21_calibration.npz":
                raise OSError("restore failed")
            return replace(path, destination)

        with (
            self._fail_second_archive(), mock.patch.object(Path, "replace", replace_file),
            self.assertRaisesRegex(RuntimeError, "rollback was incomplete; previous files are retained") as error,
        ):
            self._main("--force")
        backups = list(self.root.glob(".sets-*/s21_calibration.npz"))
        self.assertEqual(len(backups), 1)
        self.assertEqual(backups[0].read_bytes(), previous_archive)
        self.assertIn(str(backups[0].parent), str(error.exception))

    def test_native_single_is_rejected_before_any_write(self) -> None:
        self.config.radar.driver_mode = "native"
        with self.assertRaisesRegex(RuntimeError, "not supported for this radar configuration"):
            self._main()
        self.assertEqual(self._files(), {})