"""Collect complete synchronized sweep cycles from already-running devices."""

from __future__ import annotations

from collections.abc import Callable, Sequence
import queue
import threading
import time

import numpy as np

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,
    ParsedVnaDatapoint,
    find_receiver_value,
    parse_vna_datapoint_payload,
)
from python_app.hardware_full.librevna_multi_device_driver.transport import LibreVnaUsbBulkConnection

LIBREVNA_NATIVE_SWEEP_TIMEOUT_SECONDS = 1.5


def collect_complete_running_sweep_cycles(
    *,
    master_device_connection: LibreVnaUsbBulkConnection,
    slave_device_connections: Sequence[LibreVnaUsbBulkConnection],
    active_sweep_configuration: SweepConfiguration,
    cycle_count: int,
    master_stimulus_ports: Sequence[int],
    datapoint_timeout_seconds: float | None = None,
) -> SweepMeasurementResult:
    """Collect complete cycles from an already-running synchronized sweep."""
    cycle_count = int(cycle_count)
    if cycle_count < 1:
        raise ValueError("cycle_count must be >= 1")
    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")
    stage_by_master_port = {port: stage for stage, port in enumerate(stimulus_ports)}

    all_device_connections = [master_device_connection, *slave_device_connections]
    point_count = active_sweep_configuration.points
    frequencies_hz = np.zeros(point_count, dtype=np.float64)
    master_reference_measurements_by_port = {
        port: np.full((cycle_count, point_count), np.nan + 1j * np.nan, dtype=complex)
        for port in stimulus_ports
    }
    raw_receiver_measurements_by_s_parameter = {
        name: np.full((cycle_count, point_count), np.nan + 1j * np.nan, dtype=complex)
        for name in build_forward_s_parameter_names(len(slave_device_connections), stimulus_ports)
    }

    collection_errors: list[Exception] = []
    datapoint_counts_by_device_serial = {
        device_connection.serial_number: 0
        for device_connection in all_device_connections
    }
    stop_collection_requested = threading.Event()
    if datapoint_timeout_seconds is None:
        datapoint_timeout_seconds = LIBREVNA_NATIVE_SWEEP_TIMEOUT_SECONDS
    else:
        datapoint_timeout_seconds = max(0.5, float(datapoint_timeout_seconds))

    # Upper bound on how long we wait for the first cycle-start datapoint
    # (point_index==0). Without it, a device that keeps streaming non-zero
    # indices but never wraps (e.g. after a misconfigured sweep restart)
    # would refresh `last_datapoint_timestamp` on every incoming packet and
    # stall capture indefinitely.
    cycle_start_guard_seconds = max(2.0, datapoint_timeout_seconds * 4.0)

    def collect_datapoints_from_device(
        device_connection: LibreVnaUsbBulkConnection,
        handle_datapoint: Callable[[ParsedVnaDatapoint], bool],
    ) -> None:
        datapoints_received = 0
        expected_datapoint_count = cycle_count * point_count
        last_datapoint_timestamp = time.monotonic()
        collection_loop_start = last_datapoint_timestamp
        has_consumed_any_datapoint = False

        while datapoints_received < expected_datapoint_count:
            if stop_collection_requested.is_set():
                return

            now = time.monotonic()
            remaining_timeout_seconds = (last_datapoint_timestamp + datapoint_timeout_seconds) - now
            if remaining_timeout_seconds <= 0:
                collection_errors.append(
                    TimeoutError(
                        f"No datapoints from {device_connection.serial_number} for "
                        f"{datapoint_timeout_seconds:.1f} s "
                        f"(received {datapoints_received}/{point_count})"
                    )
                )
                stop_collection_requested.set()
                return

            if not has_consumed_any_datapoint and (now - collection_loop_start) > cycle_start_guard_seconds:
                collection_errors.append(
                    TimeoutError(
                        f"Device {device_connection.serial_number} streamed datapoints but never "
                        f"reached point_index=0 within {cycle_start_guard_seconds:.1f} s "
                        f"(sweep cycle did not restart)"
                    )
                )
                stop_collection_requested.set()
                return

            try:
                packet_type, payload = device_connection.receive_packet(
                    timeout_seconds=min(1.0, remaining_timeout_seconds)
                )
            except (queue.Empty, TimeoutError) as exc:
                if stop_collection_requested.is_set():
                    return
                if isinstance(exc, queue.Empty):
                    continue
                collection_errors.append(exc)
                stop_collection_requested.set()
                return
            except Exception as exc:  # noqa: BLE001
                collection_errors.append(exc)
                stop_collection_requested.set()
                return

            if packet_type != PacketType.VNA_DATAPOINT:
                continue

            parsed_datapoint = parse_vna_datapoint_payload(payload)
            if parsed_datapoint and 0 <= parsed_datapoint.point_index < point_count:
                last_datapoint_timestamp = time.monotonic()
                datapoint_was_consumed = handle_datapoint(parsed_datapoint)
                if datapoint_was_consumed:
                    has_consumed_any_datapoint = True
                    datapoint_counts_by_device_serial[device_connection.serial_number] += 1
                    datapoints_received += 1

    def build_cycle_tracking_handler(
        cycle_aware_handler: Callable[[ParsedVnaDatapoint, int], None],
    ) -> Callable[[ParsedVnaDatapoint], bool]:
        # The controller restarts the sweep before every collection, so the
        # first packet each device emits is point 0 of a brand-new cycle 0.
        # Anchoring cycle 0 on the first observed point_index==0 — instead of
        # synthesising it from a wrap — pins master and slave threads to the
        # same physical cycle even if a stale straggler from the just-stopped
        # sweep escaped the post-idle drain: such a straggler always carries
        # a non-zero point_index and is discarded until the genuine cycle 0
        # arrives. From that anchor, each subsequent wrap advances the cycle
        # counter normally.
        cycle_tracking_state = {
            "current_cycle_index": 0,
            "previous_point_index": -1,
            "synchronized": False,
        }

        def handle_datapoint(parsed_datapoint: ParsedVnaDatapoint) -> bool:
            current_point_index = parsed_datapoint.point_index

            if not cycle_tracking_state["synchronized"]:
                if current_point_index != 0:
                    return False
                cycle_tracking_state["synchronized"] = True
                cycle_tracking_state["previous_point_index"] = current_point_index
                cycle_aware_handler(parsed_datapoint, 0)
                return True

            if current_point_index < cycle_tracking_state["previous_point_index"]:
                cycle_tracking_state["current_cycle_index"] += 1
            cycle_tracking_state["previous_point_index"] = current_point_index
            current_cycle_index = cycle_tracking_state["current_cycle_index"]
            if current_cycle_index >= cycle_count:
                return False

            cycle_aware_handler(parsed_datapoint, current_cycle_index)
            return True

        return handle_datapoint

    def handle_master_datapoint(parsed_datapoint: ParsedVnaDatapoint, cycle_index: int) -> None:
        point_index = parsed_datapoint.point_index
        frequencies_hz[point_index] = parsed_datapoint.frequency_hz

        for master_stimulus_port, stage_index in stage_by_master_port.items():
            reference_receiver_value = find_receiver_value(
                parsed_datapoint.receiver_values_by_description_mask,
                stage_index=stage_index,
                is_reference_receiver=True,
                required_port_number=master_stimulus_port,
            )
            if reference_receiver_value is None:
                reference_receiver_value = find_receiver_value(
                    parsed_datapoint.receiver_values_by_description_mask,
                    stage_index=stage_index,
                    is_reference_receiver=True,
                )
            if reference_receiver_value is None:
                continue

            master_reference_measurements_by_port[master_stimulus_port][
                cycle_index,
                point_index,
            ] = reference_receiver_value

            receiver_port_number = master_stimulus_port
            s_parameter_name = master_reflection_s_parameter_name(master_stimulus_port)
            port_receiver_value = find_receiver_value(
                parsed_datapoint.receiver_values_by_description_mask,
                stage_index=stage_index,
                is_reference_receiver=False,
                required_port_number=receiver_port_number,
            )
            if port_receiver_value is not None:
                raw_receiver_measurements_by_s_parameter[s_parameter_name][
                    cycle_index,
                    point_index,
                ] = port_receiver_value

    def build_slave_datapoint_handler(slave_index: int) -> Callable[[ParsedVnaDatapoint], bool]:
        receiver_base_port = 2 * slave_index + 3

        def handle_slave_datapoint(parsed_datapoint: ParsedVnaDatapoint, cycle_index: int) -> None:
            point_index = parsed_datapoint.point_index
            for master_stimulus_port, stage_index in stage_by_master_port.items():
                first_s_parameter_name = f"S{receiver_base_port}{master_stimulus_port}"
                second_s_parameter_name = f"S{receiver_base_port + 1}{master_stimulus_port}"
                for receiver_port_number, s_parameter_name in (
                    (1, first_s_parameter_name),
                    (2, second_s_parameter_name),
                ):
                    port_receiver_value = find_receiver_value(
                        parsed_datapoint.receiver_values_by_description_mask,
                        stage_index=stage_index,
                        is_reference_receiver=False,
                        required_port_number=receiver_port_number,
                    )
                    if port_receiver_value is not None:
                        raw_receiver_measurements_by_s_parameter[s_parameter_name][
                            cycle_index,
                            point_index,
                        ] = port_receiver_value

        return build_cycle_tracking_handler(handle_slave_datapoint)

    collection_threads = [
        threading.Thread(
            target=collect_datapoints_from_device,
            args=(master_device_connection, build_cycle_tracking_handler(handle_master_datapoint)),
            daemon=True,
            name="collect-master",
        )
    ]
    for slave_index, slave_device_connection in enumerate(slave_device_connections):
        collection_threads.append(
            threading.Thread(
                target=collect_datapoints_from_device,
                args=(slave_device_connection, build_slave_datapoint_handler(slave_index)),
                daemon=True,
                name=f"collect-slave{slave_index}",
            )
        )

    for collection_thread in collection_threads:
        collection_thread.start()
    for collection_thread in collection_threads:
        collection_thread.join()

    if collection_errors:
        raise RuntimeError(f"Sweep collection failed: {collection_errors[0]}") from collection_errors[0]

    if slave_device_connections and min(datapoint_counts_by_device_serial.values(), default=0) == 0:
        raise RuntimeError(
            "No datapoints received from at least one device; hardware trigger sync did not start. "
            "Check Trigger Out/In loop and 10 MHz reference wiring."
        )

    return SweepMeasurementResult(
        frequencies_hz=frequencies_hz,
        s_parameters=calculate_last_cycle_s_parameters(
            raw_receiver_measurements_by_s_parameter,
            master_reference_measurements_by_port,
        ),
    )


def master_reflection_s_parameter_name(master_stimulus_port: int) -> str:
    """Return master reflection trace name for active output port."""
    if master_stimulus_port == 1:
        return "S11"
    return "S22"


def build_forward_s_parameter_names(slave_device_count: int, master_stimulus_ports: Sequence[int]) -> list[str]:
    """Return S-parameter names for the configured forward multi-device topology."""
    names: list[str] = []
    for master_stimulus_port in master_stimulus_ports:
        names.append(master_reflection_s_parameter_name(master_stimulus_port))
        for slave_index in range(slave_device_count):
            receiver_base_port = 2 * slave_index + 3
            names.extend(
                [
                    f"S{receiver_base_port}{master_stimulus_port}",
                    f"S{receiver_base_port + 1}{master_stimulus_port}",
                ]
            )
    return names


def calculate_last_cycle_s_parameters(
    raw_receiver_measurements_by_s_parameter: dict[str, np.ndarray],
    master_reference_measurements_by_port: dict[int, np.ndarray],
) -> dict[str, np.ndarray]:
    """Convert raw receiver captures from the final sweep cycle into S-parameters."""
    s_parameters: dict[str, np.ndarray] = {}
    for s_parameter_name, raw_receiver_measurements in raw_receiver_measurements_by_s_parameter.items():
        master_stimulus_port = int(s_parameter_name[-1])
        master_reference_measurements = master_reference_measurements_by_port[master_stimulus_port]
        if np.isnan(master_reference_measurements).any():
            missing_reference_count = int(np.isnan(master_reference_measurements).sum())
            raise RuntimeError(
                f"Master port {master_stimulus_port} reference missing for {missing_reference_count} datapoints"
            )
        if np.isnan(raw_receiver_measurements).any():
            missing_measurement_count = int(np.isnan(raw_receiver_measurements).sum())
            raise RuntimeError(
                f"Measurement {s_parameter_name} missing for {missing_measurement_count} datapoints"
            )
        s_parameters[s_parameter_name.lower()] = raw_receiver_measurements[-1] / master_reference_measurements[-1]
    return s_parameters
