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

from __future__ import annotations

from dataclasses import dataclass, field
import logging
import math
import threading
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__)

# Delays applied between successive USB reopen attempts inside recover(). Picked
# to give libusb time to re-enumerate a stuck device while staying short enough
# that a healthy reconnect feels instant. The total worst-case wait is the sum
# of all entries (1.75 s today) plus the cost of close()/open() themselves.
_REOPEN_BACKOFF_SECONDS: tuple[float, ...] = (0.25, 0.5, 1.0)

# Hard ceiling on the wall-clock time a single acquire_collection() may spend in
# the recovery path (backoff sleeps + close()/open() cost across every retry).
# Bounding this keeps a shutdown that arrives mid-recovery comfortably under the
# systemd TimeoutStopSec so the unit is never SIGKILLed for hanging on exit.
_MAX_RECOVERY_WALL_SECONDS: float = 10.0

_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 as exc:
            # Never silently latch to synthetic data: a deployed appliance must wait
            # for the real device, not record fakes. Synthetic data requires an
            # explicit backend_mode='mock' (selected in __post_init__); both 'auto'
            # and 'native' re-raise so the producer's wait-for-device retry keeps
            # trying until the hardware appears.
            logger.warning("Multi-device open failed (backend_mode=%s): %s", self.backend_mode, exc)
            raise

    def close(self) -> None:
        """Close native device transports; never raises.

        Recovery loops rely on `close()` being safe to call on a half-open or
        already-broken controller. We swallow any transport-level exception here
        and just drop the reference so the next `open()` starts fresh.
        """
        controller = self._controller
        self._controller = None
        if controller is None:
            return
        try:
            controller.close()
        except Exception as exc:  # noqa: BLE001 — recovery path, never propagate
            logger.warning("Multi-device close() ignored transport error: %s", exc)

    def recover(
        self,
        *,
        stop_event: threading.Event | None = None,
        deadline_monotonic: float | None = None,
    ) -> None:
        """Reopen native device transports after a failed acquisition.

        Tries several short backoffs so a transient USB stall does not kill the
        producer on the very first retry. Raises the last error only after
        every attempt failed — the outer acquisition loop is expected to count
        these as recovery_attempts.

        When `stop_event` is supplied the backoff waits on it instead of
        sleeping, so a shutdown request aborts the loop immediately; an optional
        `deadline_monotonic` caps the total wall-time spent here.
        """
        if self._using_mock_backend:
            return
        self.close()

        last_error: Exception | None = None
        for attempt_index, delay_s in enumerate(_REOPEN_BACKOFF_SECONDS, start=1):
            # Bail out the instant a stop is requested or the recovery budget is
            # spent, rather than committing to another (re)open attempt.
            if stop_event is not None and stop_event.is_set():
                logger.info("Multi-device recover() aborted: stop requested")
                return
            if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
                logger.warning("Multi-device recover() aborted: recovery time budget exhausted")
                break
            # Interruptible backoff: wait() returns early the moment stop is set.
            if stop_event is not None:
                if stop_event.wait(delay_s):
                    logger.info("Multi-device recover() aborted: stop requested")
                    return
            else:
                time.sleep(delay_s)
            try:
                self.open()
                if self._controller is not None:
                    logger.info(
                        "Multi-device reopen succeeded on attempt %d/%d (after %.2fs)",
                        attempt_index,
                        len(_REOPEN_BACKOFF_SECONDS),
                        delay_s,
                    )
                    return
            except Exception as exc:  # noqa: BLE001 — propagate only the last failure
                last_error = exc
                logger.warning(
                    "Multi-device reopen attempt %d/%d failed after %.2fs: %s",
                    attempt_index,
                    len(_REOPEN_BACKOFF_SECONDS),
                    delay_s,
                    exc,
                )
                self.close()  # tidy partially-opened state before next try

        if last_error is not None:
            raise last_error
        raise RuntimeError("Multi-device recover() exhausted all reopen attempts")

    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,
        *,
        stop_event: threading.Event | None = None,
    ) -> SweepCollection:
        """Acquire one complete virtual 2x4 matrix in canonical combo order.

        When `stop_event` is supplied it is threaded into the recovery path so a
        shutdown request aborts the backoff/retry loops promptly (default `None`
        preserves the original blocking behavior).
        """
        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, stop_event=stop_event
            )
        collection.capture_end_ns = time.monotonic_ns()
        return collection

    def _acquire_native_collection_with_recovery(
        self,
        collection_id: int,
        capture_start_ns: int,
        *,
        stop_event: threading.Event | None = None,
    ) -> SweepCollection:
        last_error: Exception | None = None
        # Cap total recovery wall-time across all retries so a shutdown that
        # lands mid-recovery stays well under the systemd TimeoutStopSec.
        recovery_deadline = time.monotonic() + _MAX_RECOVERY_WALL_SECONDS
        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
                # Stop the moment shutdown is requested or the recovery budget is
                # spent — do not start another reconnect we cannot finish in time.
                if stop_event is not None and stop_event.is_set():
                    break
                if time.monotonic() >= recovery_deadline:
                    logger.warning("multi-device recovery time budget exhausted; giving up")
                    break
                logger.warning(
                    "multi-device acquisition failed, reconnecting devices (%d/%d): %s",
                    attempt_index + 1,
                    self.recovery_attempts,
                    exc,
                    exc_info=True,
                )
                # recover() may itself fail when libusb cannot re-enumerate the
                # device fast enough; treat that as the same kind of recovery
                # attempt and try again on the next loop iteration, so a
                # transient USB hiccup cannot kill the whole producer.
                try:
                    self.recover(stop_event=stop_event, deadline_monotonic=recovery_deadline)
                except Exception as recover_exc:  # noqa: BLE001
                    last_error = recover_exc
                    logger.warning(
                        "multi-device recover() failed (%d/%d): %s",
                        attempt_index + 1,
                        self.recovery_attempts,
                        recover_exc,
                        exc_info=True,
                    )
                # If recovery was interrupted by a stop request, do not loop back
                # for another acquisition attempt; let shutdown proceed.
                if stop_event is not None and stop_event.is_set():
                    break

        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