"""Controller for synchronized multi-device LibreVNA sweeps."""

from __future__ import annotations

from collections.abc import Iterator, Sequence
from dataclasses import replace
from typing import Optional
import threading
import time

from python_app.hardware_full.librevna_multi_device_driver.cycle_collection import (
    collect_complete_running_sweep_cycles,
)
from python_app.hardware_full.librevna_multi_device_driver.models import (
    SweepConfiguration,
    SweepMeasurementResult,
)
from python_app.hardware_full.librevna_multi_device_driver.protocol import (
    PacketType,
    build_reference_settings_payload,
    build_sweep_settings_payload,
)
from python_app.hardware_full.librevna_multi_device_driver.transport import LibreVnaUsbBulkConnection


class MultiDeviceVnaController:
    """Coordinate one master LibreVNA and receiver slave LibreVNAs."""

    def __init__(
        self,
        master_serial_number: str,
        slave_serial_numbers: Sequence[str] = (),
        force_external_reference: bool = True,
    ) -> None:
        """Open all configured devices and initialize controller state."""
        self._master_device: LibreVnaUsbBulkConnection | None = None
        self._slave_devices: list[LibreVnaUsbBulkConnection] = []
        self._all_devices: list[LibreVnaUsbBulkConnection] = []
        self._synchronization_enabled = bool(slave_serial_numbers)
        self._force_external_reference = bool(force_external_reference)
        self._reference_configuration_applied = False
        self._last_applied_sweep_configuration: Optional[SweepConfiguration] = None
        self._last_master_stimulus_ports: tuple[int, ...] | None = None
        self._reconfigure_delay_s = 0.005
        self._sweep_is_running = False
        self._is_closed = False

        try:
            self._master_device = LibreVnaUsbBulkConnection(master_serial_number)
            self._slave_devices = [
                LibreVnaUsbBulkConnection(slave_serial_number)
                for slave_serial_number in slave_serial_numbers
            ]
            self._all_devices = [self._master_device, *self._slave_devices]
        except Exception:
            self.close()
            raise

    def __enter__(self) -> MultiDeviceVnaController:
        """Return this controller as a context manager resource."""
        return self

    def __exit__(self, *_args: object) -> None:
        """Close all devices when leaving a context manager scope."""
        self.close()

    def close(self) -> None:
        """Stop sweeping and close every opened device transport."""
        if self._is_closed:
            return

        self._is_closed = True
        self._send_idle_to_all_devices()
        for device_connection in self._all_devices:
            device_connection.close()

    def stop_continuous_sweep(self) -> None:
        """Stop the currently running sweep without closing device transports."""
        if self._is_closed:
            return
        self._send_idle_to_all_devices()

    def configure_continuous_sweep(
        self,
        sweep_configuration: SweepConfiguration,
        *,
        master_stimulus_ports: Sequence[int] = (1, 2),
    ) -> None:
        """Apply reference/sweep settings and leave devices sweeping.

        When the requested configuration already matches the running sweep,
        the host-side packet queues are drained and the device sweep is
        left untouched, so back-to-back acquires do not pay the SET_IDLE +
        SWEEP_SETTINGS round-trip cost. The drain is performed in parallel
        across devices to keep the cross-device timing skew below the USB
        latency variance, which is what previously let a hardware cycle
        wrap slip between master and slave drains and desynchronise the
        per-device cycle counters.
        """
        if self._is_closed:
            raise RuntimeError("Controller is already closed")
        stimulus_ports = self._normalize_master_stimulus_ports(master_stimulus_ports)

        if not self._reference_configuration_applied:
            self._configure_reference_clocks()

        self._drain_all_received_packets()

        if (
            self._sweep_is_running
            and self._last_applied_sweep_configuration == sweep_configuration
            and self._last_master_stimulus_ports == stimulus_ports
        ):
            return

        if self._sweep_is_running:
            self._send_idle_to_all_devices()
            time.sleep(self._reconfigure_delay_s)
            # The old sweep keeps streaming datapoints until each device
            # processes SET_IDLE. Drain again after the idle settling delay
            # so the new sweep starts on an empty queue.
            self._drain_all_received_packets()

        self._configure_sweep_on_all_devices(
            sweep_configuration,
            master_stimulus_ports=stimulus_ports,
        )

    def collect_running_sweep_cycles(
        self,
        cycle_count: int = 1,
        *,
        datapoint_timeout_seconds: float | None = None,
    ) -> SweepMeasurementResult:
        """Collect complete cycles from the already-running synchronized sweep."""
        if self._is_closed:
            raise RuntimeError("Controller is already closed")
        if (
            not self._sweep_is_running
            or self._last_applied_sweep_configuration is None
            or self._last_master_stimulus_ports is None
            or self._master_device is None
        ):
            raise RuntimeError("No running sweep is configured. Call configure_continuous_sweep() first.")

        try:
            return collect_complete_running_sweep_cycles(
                master_device_connection=self._master_device,
                slave_device_connections=self._slave_devices,
                active_sweep_configuration=self._last_applied_sweep_configuration,
                cycle_count=cycle_count,
                master_stimulus_ports=self._last_master_stimulus_ports,
                datapoint_timeout_seconds=datapoint_timeout_seconds,
            )
        except Exception:
            self._send_idle_to_all_devices()
            raise

    def iter_running_sweep_cycles(self) -> Iterator[SweepMeasurementResult]:
        """Yield complete sweep cycles forever from the configured sweep."""
        while True:
            yield self.collect_running_sweep_cycles(1)

    def _send_command_and_wait_for_acknowledgement(
        self,
        device_connection: LibreVnaUsbBulkConnection,
        packet_type: int,
        payload: bytes = b"",
        timeout_seconds: float = 3.0,
        retry_count: int = 1,
    ) -> None:
        last_error: Exception | None = None
        for _attempt_index in range(retry_count + 1):
            device_connection.send_packet(packet_type, payload)
            try:
                device_connection.wait_for_acknowledgement(timeout_seconds=timeout_seconds)
                return
            except Exception as exc:  # noqa: BLE001
                last_error = exc

        assert last_error is not None
        raise last_error

    def _try_send_command_without_failing(
        self,
        device_connection: LibreVnaUsbBulkConnection,
        packet_type: int,
        payload: bytes = b"",
        timeout_seconds: float = 3.0,
        retry_count: int = 1,
    ) -> None:
        try:
            self._send_command_and_wait_for_acknowledgement(
                device_connection,
                packet_type,
                payload,
                timeout_seconds=timeout_seconds,
                retry_count=retry_count,
            )
        except Exception:
            pass

    def _send_idle_to_all_devices(self) -> None:
        # SET_IDLE is a one-shot stop command. The ACK may be delayed only by the
        # in-flight datapoint queue, which drains within a few hundred ms. A short,
        # single-shot timeout keeps recovery snappy when one device stops responding
        # so callers (e.g. multi-radar capture) do not block for minutes on retries.
        for device_connection in self._all_devices:
            self._try_send_command_without_failing(
                device_connection,
                PacketType.SET_IDLE,
                timeout_seconds=1.0,
                retry_count=0,
            )
        self._sweep_is_running = False

    def _configure_reference_clocks(self) -> None:
        for device_connection in self._all_devices:
            self._send_command_and_wait_for_acknowledgement(
                device_connection,
                PacketType.REFERENCE_SETTINGS,
                build_reference_settings_payload(0, self._force_external_reference),
                timeout_seconds=3.0,
                retry_count=1,
            )

        time.sleep(0.05)
        self._reference_configuration_applied = True

    def _configure_sweep_on_all_devices(
        self,
        sweep_configuration: SweepConfiguration,
        *,
        master_stimulus_ports: tuple[int, ...],
    ) -> None:
        if self._master_device is None:
            raise RuntimeError("Master device is not open")

        sweep_configuration_commands = [
            *[(slave_device, False) for slave_device in self._slave_devices],
            (self._master_device, True),
        ]
        for device_connection, is_synchronization_master in sweep_configuration_commands:
            self._send_command_and_wait_for_acknowledgement(
                device_connection,
                PacketType.SWEEP_SETTINGS,
                build_sweep_settings_payload(
                    sweep_configuration,
                    is_synchronization_master=is_synchronization_master,
                    synchronization_enabled=self._synchronization_enabled,
                    master_stimulus_ports=master_stimulus_ports,
                ),
                timeout_seconds=3.0,
                retry_count=1,
            )
        self._last_applied_sweep_configuration = replace(sweep_configuration)
        self._last_master_stimulus_ports = master_stimulus_ports
        self._sweep_is_running = True

    def _drain_all_received_packets(self) -> None:
        # Drain every device queue in parallel rather than one after another:
        # serial drain leaves up to a few hundred microseconds of skew between
        # the master and slave drain moments, which is enough room for a
        # hardware cycle wrap to slip between drains and desynchronise the
        # per-device cycle counters. Each device has its own queue and lock,
        # so concurrent get_nowait calls do not contend. A single device case
        # just runs inline to avoid the thread-spawn overhead.
        if len(self._all_devices) < 2:
            for device_connection in self._all_devices:
                device_connection.drain_received_packets()
            return

        drain_threads = [
            threading.Thread(
                target=device_connection.drain_received_packets,
                name=f"drain-{device_connection.serial_number}",
                daemon=True,
            )
            for device_connection in self._all_devices
        ]
        for drain_thread in drain_threads:
            drain_thread.start()
        for drain_thread in drain_threads:
            drain_thread.join()

    @staticmethod
    def _normalize_master_stimulus_ports(master_stimulus_ports: Sequence[int]) -> tuple[int, ...]:
        stimulus_ports = tuple(int(port) for port in master_stimulus_ports)
        if not stimulus_ports:
            raise ValueError("master_stimulus_ports must not be empty")
        if len(stimulus_ports) > 2 or set(stimulus_ports) - {1, 2}:
            raise ValueError("master_stimulus_ports may contain only ports 1 and 2")
        if len(set(stimulus_ports)) != len(stimulus_ports):
            raise ValueError("master_stimulus_ports must not contain duplicates")
        return stimulus_ports
