"""One-shot workflow for capturing a full calibration set."""

from __future__ import annotations

import time

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 RunConfigModel
from python_app.storage.npz_store import NpzStore, radar_key_from_config


def capture_calibration_set(
    config: RunConfigModel,
    set_name: str,
    store: NpzStore,
) -> tuple[str, SweepCollection]:
    """Capture all switch combinations and persist them as calibration set."""
    if config.is_multi_device:
        raise RuntimeError(
            "LibreVNA multi-device S21 through calibration is not supported by this one-shot full-set helper. "
            "Use the sequential preprocess capture flow so each virtual combo can be connected through "
            "and captured explicitly."
        )

    combos = RunConfigModel.build_full_combos(config.input_switch.positions, config.output_switch.positions)

    radar = create_single_radar_service(config)
    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,
    )
    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,
    )

    traces: list[TraceData] = []
    try:
        radar.open()
        radar.configure(config.radar.sweep)
        input_switch.open()
        output_switch.open()

        for combo in combos:
            output_switch.switch_to(combo.output)
            input_switch.switch_to(combo.input)
            if config.runtime.settling_ms > 0:
                time.sleep(config.runtime.settling_ms / 1000.0)

            sweep = radar.acquire()
            traces.append(
                TraceData(
                    combo=ComboKey(input_pos=combo.input, output_pos=combo.output),
                    frequency_hz=sweep.x,
                    s11=sweep.trace("s11"),
                    s21=sweep.trace("s21"),
                )
            )
    finally:
        output_switch.close()
        input_switch.close()
        radar.close()

    collection = SweepCollection(collection_id=1, monotonic_ns=time.monotonic_ns(), traces=traces)
    radar_key = radar_key_from_config(
        model_name=config.radar.model,
        serial=config.radar.serial,
        sweep_start_hz=config.radar.sweep.start_hz,
        sweep_stop_hz=config.radar.sweep.stop_hz,
        sweep_points=config.radar.sweep.points,
        ifbw_hz=config.radar.sweep.if_bandwidth_hz,
        power_dbm=config.radar.sweep.power_dbm,
    )
    store.save_set("s21_calibration", radar_key, set_name, collection)
    return radar_key, collection