"""Sequential capture workflow for calibration/reference dataset creation."""

from __future__ import annotations

from contextlib import suppress
from dataclasses import dataclass
import time

from python_app.hardware_full.librevna_service import LibreVnaService
from python_app.hardware_full.switch_service import SwitchService
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.models.run_config_model import ComboModel, RunConfigModel
from python_app.storage.npz_store import NpzStore, radar_key_from_config


@dataclass(slots=True)
class SequentialCaptureState:
    """Immutable view of sequential capture progress."""

    kind: str
    set_name: str
    captured_count: int
    total_count: int
    current_combo: ComboModel | None


class SequentialCaptureSession:
    """Manage hardware and switch stepping for full combo capture sequence."""

    def __init__(self, config: RunConfigModel, kind: str, set_name: str) -> None:
        """Create capture session for a calibration or reference set."""
        if kind not in {"calibration", "reference"}:
            raise RuntimeError(f"Unsupported capture kind: {kind}")
        if not set_name:
            raise RuntimeError("Set name is required")

        self._config = config
        self._kind = kind
        self._set_name = set_name
        self._combos = RunConfigModel.build_full_combos(config.input_switch.positions, config.output_switch.positions)
        if not self._combos:
            raise RuntimeError("No switch combinations available for capture")

        self._traces: list[TraceData] = []
        self._next_index = 0
        self._opened = False

        self._radar = LibreVnaService(serial=config.radar.serial or None)
        self._input_switch = SwitchService(
            name=config.input_switch.name,
            positions=config.input_switch.positions,
            mode=config.input_switch.driver_mode,
            driver=config.input_switch.driver,
            gpio_chip=config.input_switch.gpio_chip,
            pin_a=config.input_switch.pin_a,
            pin_b=config.input_switch.pin_b,
            invert_logic=config.input_switch.invert_logic,
        )
        self._output_switch = SwitchService(
            name=config.output_switch.name,
            positions=config.output_switch.positions,
            mode=config.output_switch.driver_mode,
            driver=config.output_switch.driver,
            gpio_chip=config.output_switch.gpio_chip,
            pin_a=config.output_switch.pin_a,
            pin_b=config.output_switch.pin_b,
            invert_logic=config.output_switch.invert_logic,
        )

    @property
    def kind(self) -> str:
        """Return capture kind (`calibration` or `reference`)."""
        return self._kind

    @property
    def set_name(self) -> str:
        """Return destination set name."""
        return self._set_name

    def open(self) -> None:
        """Open radar and switch resources."""
        if self._opened:
            return
        self._opened = True
        try:
            self._radar.open()
            self._radar.configure(self._config.radar.sweep)
            self._input_switch.open()
            self._output_switch.open()
        except Exception:
            self.close()
            raise

    def close(self) -> None:
        """Close all opened hardware resources."""
        with suppress(Exception):
            self._output_switch.close()
        with suppress(Exception):
            self._input_switch.close()
        with suppress(Exception):
            self._radar.close()
        self._opened = False

    def state(self) -> SequentialCaptureState:
        """Return current progress snapshot."""
        current_combo = self._current_combo()
        return SequentialCaptureState(
            kind=self._kind,
            set_name=self._set_name,
            captured_count=len(self._traces),
            total_count=len(self._combos),
            current_combo=current_combo,
        )

    def capture_current_combo(self) -> TraceData:
        """Capture one trace for current combo and advance sequence cursor."""
        if not self._opened:
            raise RuntimeError("Capture session is not opened")
        combo = self._current_combo()
        if combo is None:
            raise RuntimeError("Capture session is already complete")

        self._output_switch.switch_to(combo.output)
        self._input_switch.switch_to(combo.input)
        if self._config.runtime.settling_ms > 0:
            time.sleep(self._config.runtime.settling_ms / 1000.0)

        frequency_hz, s21 = self._radar.acquire_s21()
        trace = TraceData(
            combo=ComboKey(input_pos=combo.input, output_pos=combo.output),
            frequency_hz=frequency_hz,
            s21=s21,
        )
        self._traces.append(trace)
        self._next_index += 1
        return trace

    def is_complete(self) -> bool:
        """Return `True` when all combos were captured."""
        return self._next_index >= len(self._combos)

    def finalize(self, store: NpzStore) -> tuple[str, SweepCollection]:
        """Persist completed capture into store and return radar key + collection."""
        if not self.is_complete():
            raise RuntimeError("Capture session is not complete")

        collection = SweepCollection(
            collection_id=1,
            monotonic_ns=time.monotonic_ns(),
            traces=list(self._traces),
        )
        radar_key = radar_key_from_config(
            model_name=self._config.radar.model,
            serial=self._config.radar.serial,
            sweep_start_hz=self._config.radar.sweep.start_hz,
            sweep_stop_hz=self._config.radar.sweep.stop_hz,
            sweep_points=self._config.radar.sweep.points,
            ifbw_hz=self._config.radar.sweep.if_bandwidth_hz,
            power_dbm=self._config.radar.sweep.power_dbm,
        )
        store.save_set(self._kind, radar_key, self._set_name, collection)
        return radar_key, collection

    def _current_combo(self) -> ComboModel | None:
        """Return next combo to capture, or `None` if session is complete."""
        if self._next_index >= len(self._combos):
            return None
        return self._combos[self._next_index]
