"""High-level service for LibreVNA multi-device acquisition."""

from __future__ import annotations

from dataclasses import dataclass, field
import logging
import math
import time
from typing import TYPE_CHECKING

import numpy as np

from python_app.hardware_full.librevna_multi_device_driver.cycle_collection import LIBREVNA_NATIVE_SWEEP_TIMEOUT_SECONDS
from python_app.hardware_full.librevna_multi_device_driver.models import SweepConfiguration
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.models.run_config_model import RadarSweepModel, RunConfigModel

if TYPE_CHECKING:
    from python_app.hardware_full.librevna_multi_device_driver.controller import MultiDeviceVnaController

logger = logging.getLogger(__name__)

_INPUT_S_PARAMETERS_BY_OUTPUT: dict[int, tuple[str, ...]] = {
    0: ("s31", "s41", "s51", "s61"),
    1: ("s32", "s42", "s52", "s62"),
}

@dataclass(slots=True)
class MultiDeviceLibreVnaService:
    """Acquire a virtual 2x4 switch matrix from synchronized LibreVNA devices."""

    master_serial: str
    slave_serials: list[str]
    force_external_reference: bool = True
    recovery_attempts: int = 3
    backend_mode: str = "auto"
    _controller: "MultiDeviceVnaController | None" = field(init=False, default=None, repr=False)
    _sweep_configuration: SweepConfiguration | None = field(init=False, default=None, repr=False)
    _using_mock_backend: bool = field(init=False, default=False, repr=False)
    _mock_phase: float = field(init=False, default=0.0, repr=False)

    def __post_init__(self) -> None:
        """Validate static topology and backend selection."""
        self.master_serial = str(self.master_serial).strip()
        self.slave_serials = [str(value).strip() for value in self.slave_serials if str(value).strip()]
        if len(self.slave_serials) != 2:
            raise ValueError("LibreVNA multi-device mode requires exactly two slave serial numbers")
        self.recovery_attempts = max(0, int(self.recovery_attempts))

        mode = self.backend_mode.strip().lower()
        if mode not in {"auto", "native", "mock"}:
            raise ValueError(f"Unsupported multi-device backend mode: {self.backend_mode}")
        self.backend_mode = mode
        self._using_mock_backend = mode == "mock"

    @property
    def using_mock_backend(self) -> bool:
        """Return whether this service is generating synthetic data."""
        return self._using_mock_backend

    def open(self) -> None:
        """Open native device transports when not in mock mode."""
        if self._using_mock_backend or self._controller is not None:
            return
        try:
            from python_app.hardware_full.librevna_multi_device_driver.controller import MultiDeviceVnaController

            self._controller = MultiDeviceVnaController(
                master_serial_number=self.master_serial,
                slave_serial_numbers=self.slave_serials,
                force_external_reference=self.force_external_reference,
            )
        except Exception:
            if self.backend_mode == "native":
                raise
            self._using_mock_backend = True
            self._controller = None

    def close(self) -> None:
        """Close native device transports."""
        if self._controller is not None:
            self._controller.close()
            self._controller = None

    def recover(self) -> None:
        """Reopen native device transports after a failed acquisition."""
        if self._using_mock_backend:
            return
        self.close()
        time.sleep(0.25)
        self.open()

    def configure(self, sweep: RadarSweepModel) -> None:
        """Store sweep settings for subsequent full-matrix acquisitions."""
        self._sweep_configuration = SweepConfiguration(
            start_hz=int(round(float(sweep.start_hz))),
            stop_hz=int(round(float(sweep.stop_hz))),
            points=int(sweep.points),
            if_bandwidth=int(round(float(sweep.if_bandwidth_hz))),
            power_dbm=float(sweep.power_dbm),
        )

    def acquire_collection(self, collection_id: int = 1) -> SweepCollection:
        """Acquire one complete virtual 2x4 matrix in canonical combo order."""
        if self._sweep_configuration is None:
            raise RuntimeError("Multi-device service is not configured")
        capture_start_ns = time.monotonic_ns()
        if self._using_mock_backend:
            collection = self._acquire_mock_collection(collection_id, capture_start_ns)
        else:
            collection = self._acquire_native_collection_with_recovery(collection_id, capture_start_ns)
        collection.capture_end_ns = time.monotonic_ns()
        return collection

    def _acquire_native_collection_with_recovery(
        self,
        collection_id: int,
        capture_start_ns: int,
    ) -> SweepCollection:
        last_error: Exception | None = None
        for attempt_index in range(self.recovery_attempts + 1):
            try:
                return self._acquire_native_collection(collection_id, capture_start_ns)
            except Exception as exc:  # noqa: BLE001
                last_error = exc
                if attempt_index >= self.recovery_attempts:
                    break
                logger.warning(
                    "multi-device acquisition failed, reconnecting devices (%d/%d): %s",
                    attempt_index + 1,
                    self.recovery_attempts,
                    exc,
                    exc_info=True,
                )
                self.recover()

        assert last_error is not None
        raise last_error

    def _acquire_native_collection(self, collection_id: int, capture_start_ns: int) -> SweepCollection:
        if self._controller is None:
            raise RuntimeError("Multi-device controller is not open")
        assert self._sweep_configuration is not None

        self._controller.configure_continuous_sweep(self._sweep_configuration)
        result = self._controller.collect_running_sweep_cycles(
            1,
            datapoint_timeout_seconds=LIBREVNA_NATIVE_SWEEP_TIMEOUT_SECONDS,
        )
        normalized_s_parameters = {
            str(name).lower(): np.asarray(values, dtype=np.complex64)
            for name, values in result.s_parameters.items()
        }
        frequencies = np.asarray(result.frequencies_hz, dtype=np.float32)

        traces: list[TraceData] = []
        for output_pos in range(RunConfigModel.MULTI_DEVICE_OUTPUT_POSITIONS):
            reflection = self._required_s_parameter(
                normalized_s_parameters,
                "s11" if output_pos == 0 else "s22",
            )

            for input_pos, s_parameter_name in enumerate(_INPUT_S_PARAMETERS_BY_OUTPUT[output_pos]):
                traces.append(
                    TraceData(
                        combo=ComboKey(input=input_pos, output=output_pos),
                        frequency_hz=frequencies,
                        s11=reflection,
                        s21=self._required_s_parameter(normalized_s_parameters, s_parameter_name),
                    )
                )

        return SweepCollection(
            collection_id=int(collection_id),
            monotonic_ns=time.monotonic_ns(),
            traces=traces,
            capture_start_ns=capture_start_ns,
        )

    def _acquire_mock_collection(self, collection_id: int, capture_start_ns: int) -> SweepCollection:
        assert self._sweep_configuration is not None
        points = int(self._sweep_configuration.points)
        frequencies = np.linspace(
            self._sweep_configuration.start_hz,
            self._sweep_configuration.stop_hz,
            points,
            dtype=np.float32,
        )
        base_phase = 2.0 * math.pi * np.linspace(0.0, 1.0, points, dtype=np.float32) + self._mock_phase
        traces: list[TraceData] = []
        for output_pos in range(RunConfigModel.MULTI_DEVICE_OUTPUT_POSITIONS):
            reflected_phase = base_phase * (0.55 + 0.05 * output_pos) + 0.7 * (output_pos + 1)
            s11 = (
                (0.22 + 0.04 * output_pos) * np.cos(reflected_phase)
                + 1j * (0.22 + 0.04 * output_pos) * np.sin(reflected_phase)
            ).astype(np.complex64)
            for input_pos in range(RunConfigModel.MULTI_DEVICE_INPUT_POSITIONS):
                gain = 0.45 + 0.08 * input_pos + 0.03 * output_pos
                phase = base_phase * (1.0 + 0.03 * input_pos) + (0.4 * input_pos + 0.9 * output_pos)
                s21 = (gain * np.cos(phase) + 1j * gain * np.sin(phase)).astype(np.complex64)
                traces.append(
                    TraceData(
                        combo=ComboKey(input=input_pos, output=output_pos),
                        frequency_hz=frequencies,
                        s11=s11,
                        s21=s21,
                    )
                )
        self._mock_phase += 0.05
        return SweepCollection(
            collection_id=int(collection_id),
            monotonic_ns=time.monotonic_ns(),
            traces=traces,
            capture_start_ns=capture_start_ns,
        )

    @staticmethod
    def _required_s_parameter(s_parameters: dict[str, np.ndarray], name: str) -> np.ndarray:
        values = s_parameters.get(name)
        if values is None:
            raise RuntimeError(f"Multi-device sweep is missing required {name.upper()} trace")
        return values