"""Sequential capture workflow for preprocess asset dataset creation."""

from __future__ import annotations

from contextlib import suppress
from dataclasses import dataclass
import time

import numpy as np

from python_app.hardware_full.multi_device_service import MultiDeviceLibreVnaService
from python_app.hardware_full.single_radar_service import create_single_radar_service
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

MULTI_DEVICE_MANUAL_CAPTURE_KINDS = frozenset({"s21_calibration", "s11_open", "s11_short", "s11_load"})


@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
    can_undo: bool
    is_complete: bool
    variant_count: int = 1
    supports_batch_capture: bool = True


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 one preprocess asset set."""
        if kind not in {"s21_calibration", "s21_reference", "s11_open", "s11_short", "s11_load", "s11_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._is_multi_device = config.is_multi_device
        self._manual_multi_device_capture = self._is_multi_device and kind in MULTI_DEVICE_MANUAL_CAPTURE_KINDS
        self._combos = (
            RunConfigModel.build_multi_device_virtual_combos()
            if self._is_multi_device
            else 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

        if self._is_multi_device:
            self._radar = MultiDeviceLibreVnaService(
                master_serial=config.radar.serial,
                slave_serials=list(config.radar.multi_device.slave_serials),
                force_external_reference=config.radar.multi_device.force_external_reference,
                recovery_attempts=config.radar.multi_device.recovery_attempts,
                backend_mode=config.radar.driver_mode,
            )
            self._input_switch = None
            self._output_switch = None
        else:
            self._radar = create_single_radar_service(config)
            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 canonical preprocess asset key for this capture session."""
        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)
            if self._input_switch is not None:
                self._input_switch.open()
            if self._output_switch is not None:
                self._output_switch.open()
        except Exception:
            self.close()
            raise

    def close(self) -> None:
        """Close all opened hardware resources."""
        with suppress(Exception):
            if self._output_switch is not None:
                self._output_switch.close()
        with suppress(Exception):
            if self._input_switch is not None:
                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,
            can_undo=bool(self._traces),
            is_complete=self.is_complete(),
            supports_batch_capture=not self._manual_multi_device_capture,
        )

    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")

        if self._is_multi_device:
            collection = self._radar.acquire_collection(collection_id=1)
            if not collection.traces:
                raise RuntimeError("Multi-device capture returned no traces")
            if self._manual_multi_device_capture:
                trace = select_trace_for_combo(collection, combo)
                self._traces.append(trace)
                self._next_index += 1
                return trace

            self._traces.extend(collection.traces)
            self._next_index = len(self._combos)
            return collection.traces[-1]

        assert self._input_switch is not None
        assert self._output_switch is not None
        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)

        sweep = self._radar.acquire()
        trace = TraceData(
            combo=ComboKey(input_pos=combo.input, output_pos=combo.output),
            frequency_hz=np.asarray(sweep.x, dtype=np.float32),
            s11=np.asarray(sweep.trace("s11"), dtype=np.complex64),
            s21=np.asarray(sweep.trace("s21"), dtype=np.complex64),
        )
        self._traces.append(trace)
        self._next_index += 1
        return trace

    def undo_last_capture(self) -> TraceData:
        """Remove the most recently captured trace and rewind cursor by one combo."""
        if not self._opened:
            raise RuntimeError("Capture session is not opened")
        if not self._traces or self._next_index <= 0:
            raise RuntimeError("No captured combo is available to undo")

        if self._is_multi_device and not self._manual_multi_device_capture:
            if len(self._traces) != len(self._combos):
                raise RuntimeError("Capture session state is inconsistent; multi-device trace matrix is incomplete")
            removed_trace = self._traces[-1]
            self._traces.clear()
            self._next_index = 0
            return removed_trace

        expected_combo = self._combos[self._next_index - 1]
        removed_trace = self._traces[-1]
        if (
            int(removed_trace.combo.input_pos) != int(expected_combo.input)
            or int(removed_trace.combo.output_pos) != int(expected_combo.output)
        ):
            raise RuntimeError("Capture session state is inconsistent; last trace does not match rewind combo")
        self._next_index -= 1
        self._traces.pop()
        return removed_trace

    def last_captured_trace(self) -> TraceData | None:
        """Return the most recently captured trace, if any."""
        if not self._traces:
            return None
        return self._traces[-1]

    def captured_traces(self) -> list[TraceData]:
        """Return captured traces in acquisition order."""
        return list(self._traces)

    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,
            extra_serials=self._config.radar_key_extra_parts() or None,
        )
        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]


def select_trace_for_combo(collection: SweepCollection, combo: ComboModel) -> TraceData:
    """Return the trace matching a virtual combo from a full multi-device capture."""
    for trace in collection.traces:
        if (
            int(trace.combo.input_pos) == int(combo.input)
            and int(trace.combo.output_pos) == int(combo.output)
        ):
            return trace
    raise RuntimeError(f"Multi-device capture is missing trace for input={combo.input}, output={combo.output}")
