"""Public datamodels and low-level packet payload representations."""

from __future__ import annotations

from dataclasses import dataclass, field
import logging
from threading import Event, Thread
from typing import Any, Callable

import numpy as np

from .enums import HardwareFamily, PacketType, SweepKind, SweepScale, SyncMode

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class Packet:
    """Generic protocol packet container."""

    type: PacketType
    payload: Any = b""


@dataclass(slots=True)
class USBDeviceDescriptor:
    """USB device descriptor exposed by transport discovery."""

    serial: str
    vendor_id: int
    product_id: int


@dataclass(slots=True)
class DeviceLimits:
    """Hardware capability limits reported by the instrument."""

    min_frequency_hz: float
    max_frequency_hz: float
    max_frequency_harmonic_hz: float
    min_ifbw_hz: float
    max_ifbw_hz: float
    max_points: int
    min_power_dbm: float
    max_power_dbm: float
    min_rbw_hz: float
    max_rbw_hz: float
    max_amplitude_points: int
    max_dwell_time_s: float


@dataclass(slots=True)
class DeviceInfo:
    """Device identity and capabilities."""

    protocol_version: int
    firmware_major: int
    firmware_minor: int
    firmware_patch: int
    firmware_version: str
    hardware_version: int
    hardware_revision: str
    hardware_family: HardwareFamily
    limits: DeviceLimits
    num_ports: int


@dataclass(slots=True)
class DeviceStatus:
    """Current runtime status and telemetry."""

    family: HardwareFamily
    source_locked: bool | None
    lo_locked: bool | None
    adc_overload: bool | None
    unlevel: bool | None
    temperatures_c: list[float] = field(default_factory=list)
    raw: dict[str, int | float | bool] = field(default_factory=dict)


@dataclass(slots=True)
class VNASweepSettings:
    """Configuration for one VNA sweep setup packet."""

    kind: SweepKind = SweepKind.FREQUENCY
    f_start_hz: float = 1_000_000.0
    f_stop_hz: float = 6_000_000_000.0
    points: int = 501
    if_bandwidth_hz: float = 1_000.0
    power_start_dbm: float = -10.0
    power_stop_dbm: float = -10.0
    excited_ports: tuple[int, ...] = (1, 2)
    sweep_scale: SweepScale = SweepScale.LIN
    dwell_s: float = 0.0
    suppress_invalid_peaks: bool = True
    fixed_power_setting: bool = False
    standby: bool = True
    sync_mode: SyncMode | None = None
    sync_master: bool = False

    def __post_init__(self) -> None:
        """Validate configuration before transmitting to the device."""
        if self.points <= 0:
            raise ValueError("points must be > 0")
        if self.if_bandwidth_hz <= 0:
            raise ValueError("if_bandwidth_hz must be > 0")
        if self.dwell_s < 0:
            raise ValueError("dwell_s must be >= 0")
        if not self.excited_ports:
            raise ValueError("excited_ports must not be empty")
        if len(set(self.excited_ports)) != len(self.excited_ports):
            raise ValueError("excited_ports must not contain duplicates")
        if any(port <= 0 for port in self.excited_ports):
            raise ValueError("excited_ports must use 1-based positive port numbers")
        if self.f_stop_hz < self.f_start_hz:
            raise ValueError("f_stop_hz must be >= f_start_hz")
        if self.kind == SweepKind.POWER and self.power_stop_dbm < self.power_start_dbm:
            raise ValueError("power_stop_dbm must be >= power_start_dbm for power sweep")
        if self.kind == SweepKind.POWER and self.f_start_hz != self.f_stop_hz:
            raise ValueError("power sweep requires f_start_hz == f_stop_hz")


@dataclass(slots=True)
class GeneratorSettings:
    """Configuration for LibreVNA signal-generator mode."""

    frequency_hz: float = 1_000_000.0
    power_dbm: float = -10.0
    active_port: int = 1
    apply_amplitude_correction: bool = False

    def __post_init__(self) -> None:
        """Validate generator settings before transmission."""
        if self.frequency_hz <= 0:
            raise ValueError("frequency_hz must be > 0")
        if self.active_port not in {1, 2}:
            raise ValueError("active_port must be 1 or 2")


@dataclass(slots=True)
class DeviceConfigVariant:
    """Family-specific device configuration fields."""

    family: HardwareFamily
    values: dict[str, int | float | bool] = field(default_factory=dict)


@dataclass(slots=True)
class VNADatapointPacket:
    """Decoded low-level VNADatapoint payload."""

    frequency_or_time: int
    cdbm: int
    point_number: int
    real: np.ndarray
    imag: np.ndarray
    flags: np.ndarray


@dataclass(slots=True)
class VNARawPoint:
    """Normalized VNA datapoint used by streaming callback API."""

    point_number: int
    frequency_hz: float | None
    time_s: float | None
    power_dbm: float | None
    measurements: dict[str, complex] = field(default_factory=dict)
    z0: float = 50.0


class StreamHandle:
    """Handle returned from streaming subscriptions."""

    def __init__(
        self,
        stop_event: Event,
        close_callback: Callable[[], None],
        reader_thread: Thread | None = None,
    ) -> None:
        """Create stream handle with stop event and unsubscribe callback."""
        self._stop_event = stop_event
        self._close_callback = close_callback
        self._reader_thread = reader_thread

    def close(self, *, join_timeout_s: float | None = 1.0) -> None:
        """Stop streaming and release internal resources."""
        logger.debug("Closing stream handle (join_timeout_s=%s)", join_timeout_s)
        if not self._stop_event.is_set():
            self._stop_event.set()
            self._close_callback()
        if self._reader_thread is not None and self._reader_thread.is_alive():
            self._reader_thread.join(timeout=join_timeout_s)
        logger.debug("Stream handle closed")


@dataclass(slots=True)
class SweepResult:
    """Container for complex VNA traces sharing one X-axis."""

    x: np.ndarray
    traces: dict[str, np.ndarray] = field(default_factory=dict)
    x_label: str = "frequency_hz"

    def __post_init__(self) -> None:
        """Normalize dtypes and validate shape compatibility."""
        self.x = np.asarray(self.x, dtype=np.float64)
        normalized: dict[str, np.ndarray] = {}
        for key, values in self.traces.items():
            arr = np.asarray(values, dtype=np.complex128)
            if arr.shape != self.x.shape:
                raise ValueError(
                    f"Trace '{key}' shape {arr.shape} does not match axis shape {self.x.shape}"
                )
            normalized[key.strip().lower()] = arr
        self.traces = normalized

    @property
    def s11(self) -> np.ndarray | None:
        """Return `S11` trace when available."""
        return self.traces.get("s11")

    @property
    def s21(self) -> np.ndarray | None:
        """Return `S21` trace when available."""
        return self.traces.get("s21")

    @property
    def s12(self) -> np.ndarray | None:
        """Return `S12` trace when available."""
        return self.traces.get("s12")

    @property
    def s22(self) -> np.ndarray | None:
        """Return `S22` trace when available."""
        return self.traces.get("s22")

    def trace(self, parameter: str) -> np.ndarray:
        """Return complex trace by parameter name."""
        key = parameter.strip().lower()
        if key not in self.traces:
            raise KeyError(f"Trace '{parameter}' is not available")
        return self.traces[key]

    def real(self, parameter: str) -> np.ndarray:
        """Return real part of selected trace."""
        return self.trace(parameter).real

    def imag(self, parameter: str) -> np.ndarray:
        """Return imaginary part of selected trace."""
        return self.trace(parameter).imag

    def to_npz(self, path: str) -> None:
        """Save result as NumPy `.npz` archive."""
        data: dict[str, np.ndarray] = {self.x_label: self.x}
        data.update(self.traces)
        np.savez(path, **data)

    def to_csv(self, path: str) -> None:
        """Save result as CSV with `<trace>_real`/`<trace>_imag` columns."""
        columns: list[np.ndarray] = [self.x]
        headers: list[str] = [self.x_label]
        for name, values in sorted(self.traces.items()):
            columns.append(values.real)
            columns.append(values.imag)
            headers.append(f"{name}_real")
            headers.append(f"{name}_imag")
        matrix = np.column_stack(columns)
        np.savetxt(path, matrix, delimiter=",", header=",".join(headers), comments="")


PacketPayload = Any