"""Sequential capture workflow that captures each combo across multiple radar sweep variants."""

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
from python_app.workflows.radar_config_variants import RadarConfigVariant
from python_app.workflows.sequential_capture_workflow import (
    MULTI_DEVICE_MANUAL_CAPTURE_KINDS,
    SequentialCaptureState,
    combine_collections_via_median,
    combine_traces_via_median,
    select_trace_for_combo,
)


@dataclass(frozen=True, slots=True)
class MultiRadarCaptureBatch:
    """One completed combo capture across all configured radar variants."""

    combo: ComboModel
    traces: tuple[TraceData, ...]
    variant_labels: tuple[str, ...]

    @property
    def display_trace(self) -> TraceData:
        """Return the last trace in the batch for preview rendering."""
        return self.traces[-1]


@dataclass(frozen=True, slots=True)
class MultiRadarSavedSet:
    """One preprocess set persisted for one radar variant."""

    display_name: str
    radar_key: str
    trace_count: int


class MultiRadarSequentialCaptureSession:
    """Capture a full combo matrix for one preprocess asset across multiple radar sweep variants."""

    def __init__(
        self,
        *,
        base_config: RunConfigModel,
        kind: str,
        set_name: str,
        radar_variants: list[RadarConfigVariant],
        median_sweep_count: int = 1,
    ) -> None:
        """Create capture session for one preprocess asset set and multiple radar variants."""
        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 not radar_variants:
            raise RuntimeError("At least one radar variant is required")
        if int(median_sweep_count) < 1:
            raise RuntimeError("median_sweep_count must be >= 1")

        self._base_config = base_config
        self._kind = kind
        self._set_name = set_name
        self._radar_variants = list(radar_variants)
        self._median_sweep_count = int(median_sweep_count)
        self._is_matrix_radar = base_config.is_matrix_radar
        self._is_multi_device = base_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(
                base_config.input_switch.positions,
                base_config.output_switch.positions,
            )
        )
        if not self._combos:
            raise RuntimeError("No switch combinations available for capture")

        self._captured_batches: list[MultiRadarCaptureBatch] = []
        self._traces_by_radar_key = {
            variant.radar_key: []
            for variant in self._radar_variants
        }
        self._next_index = 0
        self._opened = False

        if self._is_matrix_radar:
            self._radar: MatrixRadarService = create_matrix_radar_service(base_config)
            self._input_switch = None
            self._output_switch = None
        else:
            self._radar = create_single_radar_service(base_config)
            self._input_switch = SwitchService(
                name=base_config.input_switch.name,
                positions=base_config.input_switch.positions,
                mode=base_config.input_switch.driver_mode,
                driver=base_config.input_switch.driver,
                gpio_chip=base_config.input_switch.gpio_chip,
                pin_a=base_config.input_switch.pin_a,
                pin_b=base_config.input_switch.pin_b,
                invert_logic=base_config.input_switch.invert_logic,
            )
            self._output_switch = SwitchService(
                name=base_config.output_switch.name,
                positions=base_config.output_switch.positions,
                mode=base_config.output_switch.driver_mode,
                driver=base_config.output_switch.driver,
                gpio_chip=base_config.output_switch.gpio_chip,
                pin_a=base_config.output_switch.pin_a,
                pin_b=base_config.output_switch.pin_b,
                invert_logic=base_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._base_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=(
                self._next_index
                if self._is_matrix_radar and not self._manual_multi_device_capture
                else len(self._captured_batches)
            ),
            total_count=len(self._combos),
            current_combo=current_combo,
            can_undo=bool(self._captured_batches),
            is_complete=self.is_complete(),
            variant_count=len(self._radar_variants),
            supports_batch_capture=not self._manual_multi_device_capture,
        )

    def capture_current_combo(self) -> MultiRadarCaptureBatch:
        """Capture the current combo across all radar variants and advance the combo 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")

        pending_traces_by_radar_key: dict[str, list[TraceData]] = {}
        display_traces: list[TraceData] = []
        variant_labels: list[str] = []

        if self._is_matrix_radar:
            for variant in self._radar_variants:
                self._radar.configure(variant.config.radar.sweep)
                if self._base_config.runtime.settling_ms > 0:
                    time.sleep(self._base_config.runtime.settling_ms / 1000.0)
                collections: list[SweepCollection] = []
                for _ in range(self._median_sweep_count):
                    collection = self._radar.acquire_collection(collection_id=1)
                    if not collection.traces:
                        raise RuntimeError(
                            f"Matrix radar variant {variant.display_name} 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)
                    pending_traces_by_radar_key[variant.radar_key] = [trace]
                    display_traces.append(trace)
                else:
                    combined_collection = combine_collections_via_median(collections)
                    pending_traces_by_radar_key[variant.radar_key] = list(combined_collection.traces)
                    display_traces.append(combined_collection.traces[-1])
                variant_labels.append(variant.display_name)
        else:
            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._base_config.runtime.settling_ms > 0:
                time.sleep(self._base_config.runtime.settling_ms / 1000.0)

            for variant in self._radar_variants:
                self._radar.configure(variant.config.radar.sweep)
                if self._base_config.runtime.settling_ms > 0:
                    time.sleep(self._base_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)
                pending_traces_by_radar_key[variant.radar_key] = [trace]
                display_traces.append(trace)
                variant_labels.append(variant.display_name)

        for radar_key, traces in pending_traces_by_radar_key.items():
            self._traces_by_radar_key[radar_key].extend(traces)
        batch = MultiRadarCaptureBatch(
            combo=combo,
            traces=tuple(display_traces),
            variant_labels=tuple(variant_labels),
        )
        self._captured_batches.append(batch)
        if self._is_matrix_radar and not self._manual_multi_device_capture:
            self._next_index = len(self._combos)
        else:
            self._next_index += 1
        return batch

    def undo_last_capture(self) -> MultiRadarCaptureBatch:
        """Remove the most recently captured combo batch and rewind the cursor."""
        if not self._opened:
            raise RuntimeError("Capture session is not opened")
        if not self._captured_batches 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:
            removed_batch = self._captured_batches[-1]
            for variant in self._radar_variants:
                traces = self._traces_by_radar_key[variant.radar_key]
                if len(traces) < len(self._combos):
                    raise RuntimeError("Capture session state is inconsistent; missing matrix radar traces")
                del traces[-len(self._combos) :]
            self._next_index = 0
            self._captured_batches.pop()
            return removed_batch

        expected_combo = self._combos[self._next_index - 1]
        removed_batch = self._captured_batches[-1]
        if (
            int(removed_batch.combo.input) != int(expected_combo.input)
            or int(removed_batch.combo.output) != int(expected_combo.output)
        ):
            raise RuntimeError("Capture session state is inconsistent; last batch does not match rewind combo")

        for variant in self._radar_variants:
            traces = self._traces_by_radar_key[variant.radar_key]
            if not traces:
                raise RuntimeError("Capture session state is inconsistent; missing trace during undo")
            traces.pop()

        self._next_index -= 1
        self._captured_batches.pop()
        return removed_batch

    def last_captured_trace(self) -> TraceData | None:
        """Return the most recent trace from the most recent combo batch, if any."""
        if not self._captured_batches:
            return None
        return self._captured_batches[-1].display_trace

    def captured_batches(self) -> list[MultiRadarCaptureBatch]:
        """Return completed combo batches in capture order."""
        return list(self._captured_batches)

    def radar_variant_count(self) -> int:
        """Return how many radar variants are captured per combo."""
        return len(self._radar_variants)

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

    def finalize(self, store: NpzStore) -> list[MultiRadarSavedSet]:
        """Persist completed captures as one preprocess set per radar variant."""
        if not self.is_complete():
            raise RuntimeError("Capture session is not complete")

        monotonic_ns = time.monotonic_ns()
        saved_sets: list[MultiRadarSavedSet] = []
        for variant in self._radar_variants:
            traces = list(self._traces_by_radar_key[variant.radar_key])
            collection = SweepCollection(
                collection_id=1,
                monotonic_ns=monotonic_ns,
                traces=traces,
            )
            store.save_set(self._kind, variant.radar_key, self._set_name, collection)
            saved_sets.append(
                MultiRadarSavedSet(
                    display_name=variant.display_name,
                    radar_key=variant.radar_key,
                    trace_count=len(traces),
                )
            )
        return saved_sets

    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]
