"""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.matrix_radar_service import MatrixRadarService, create_matrix_radar_service
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"})
DEFAULT_CALIBRATION_MEDIAN_SWEEP_COUNT = 5


@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,
        *,
        median_sweep_count: int = 1,
    ) -> 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")
        if int(median_sweep_count) < 1:
            raise RuntimeError("median_sweep_count must be >= 1")

        self._config = config
        self._kind = kind
        self._set_name = set_name
        self._median_sweep_count = int(median_sweep_count)
        self._is_matrix_radar = config.is_matrix_radar
        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_matrix_radar_virtual_combos()
            if self._is_matrix_radar
            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_matrix_radar:
            self._radar: MatrixRadarService = create_matrix_radar_service(config)
            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_matrix_radar:
            collections: list[SweepCollection] = []
            for _ in range(self._median_sweep_count):
                collection = self._radar.acquire_collection(collection_id=1)
                if not collection.traces:
                    raise RuntimeError("Matrix radar capture returned no traces")
                collections.append(collection)
            if self._manual_multi_device_capture:
                per_sweep_traces = [select_trace_for_combo(collection, combo) for collection in collections]
                trace = combine_traces_via_median(per_sweep_traces)
                self._traces.append(trace)
                self._next_index += 1
                return trace

            combined_collection = combine_collections_via_median(collections)
            self._traces.extend(combined_collection.traces)
            self._next_index = len(self._combos)
            return combined_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_traces: list[TraceData] = []
        for _ in range(self._median_sweep_count):
            sweep = self._radar.acquire()
            sweep_traces.append(
                TraceData(
                    combo=ComboKey(input=combo.input, output=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),
                )
            )
        trace = combine_traces_via_median(sweep_traces)
        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_matrix_radar and not self._manual_multi_device_capture:
            if len(self._traces) != len(self._combos):
                raise RuntimeError("Capture session state is inconsistent; matrix radar 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) != int(expected_combo.input)
            or int(removed_trace.combo.output) != 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) == int(combo.input)
            and int(trace.combo.output) == int(combo.output)
        ):
            return trace
    raise RuntimeError(f"Multi-device capture is missing trace for input={combo.input}, output={combo.output}")


def combine_traces_via_median(traces: list[TraceData]) -> TraceData:
    """Return one trace whose S11/S21 are the per-point median of the inputs.

    A single-element input is returned unchanged. With multiple inputs, the real
    and imaginary parts of each complex sample are medianed independently so a
    single bad sweep (e.g. an outlier with random phase) is rejected without
    corrupting the saved calibration trace.
    """
    if not traces:
        raise RuntimeError("Cannot combine empty sweep list")
    if len(traces) == 1:
        return traces[0]

    first = traces[0]
    combo = first.combo
    point_count = first.frequency_hz.size
    for index, trace in enumerate(traces[1:], start=1):
        if (
            int(trace.combo.input) != int(combo.input)
            or int(trace.combo.output) != int(combo.output)
        ):
            raise RuntimeError(
                f"Median combine combo mismatch at sweep {index}: "
                f"({trace.combo.input},{trace.combo.output}) "
                f"vs ({combo.input},{combo.output})"
            )
        if trace.frequency_hz.size != point_count:
            raise RuntimeError(
                f"Median combine point-count mismatch at sweep {index}: "
                f"{trace.frequency_hz.size} vs {point_count}"
            )

    s11_stack = np.stack([np.asarray(t.s11, dtype=np.complex64) for t in traces], axis=0)
    s21_stack = np.stack([np.asarray(t.s21, dtype=np.complex64) for t in traces], axis=0)
    s11_median = (
        np.median(s11_stack.real, axis=0) + 1j * np.median(s11_stack.imag, axis=0)
    ).astype(np.complex64)
    s21_median = (
        np.median(s21_stack.real, axis=0) + 1j * np.median(s21_stack.imag, axis=0)
    ).astype(np.complex64)
    return TraceData(
        combo=ComboKey(input=int(combo.input), output=int(combo.output)),
        frequency_hz=np.asarray(first.frequency_hz, dtype=np.float32),
        s11=s11_median,
        s21=s21_median,
    )


def combine_collections_via_median(collections: list[SweepCollection]) -> SweepCollection:
    """Combine multi-device matrix captures into one collection with per-combo medians."""
    if not collections:
        raise RuntimeError("Cannot combine empty collection list")
    if len(collections) == 1:
        return collections[0]

    reference = collections[0]
    expected_combos = [(trace.combo.input, trace.combo.output) for trace in reference.traces]
    medianed_traces: list[TraceData] = []
    for combo_index, (input_pos, output_pos) in enumerate(expected_combos):
        per_sweep_traces: list[TraceData] = []
        for collection_index, collection in enumerate(collections):
            if combo_index >= len(collection.traces):
                raise RuntimeError(
                    f"Median collections have mismatched trace counts at sweep {collection_index}"
                )
            trace = collection.traces[combo_index]
            if (
                int(trace.combo.input) != int(input_pos)
                or int(trace.combo.output) != int(output_pos)
            ):
                raise RuntimeError(
                    f"Median collections trace order mismatch at sweep {collection_index}, "
                    f"combo index {combo_index}"
                )
            per_sweep_traces.append(trace)
        medianed_traces.append(combine_traces_via_median(per_sweep_traces))

    return SweepCollection(
        collection_id=int(reference.collection_id),
        monotonic_ns=int(reference.monotonic_ns),
        traces=medianed_traces,
        capture_start_ns=int(reference.capture_start_ns),
        capture_end_ns=int(reference.capture_end_ns),
    )
