"""TCP client for a Compact-M K209 connected to a remote S2VNA host."""

from __future__ import annotations

from dataclasses import dataclass, field
import socket

import numpy as np

from python_app.hardware_full.k209_remote_protocol import (
    COMMAND_ACQUIRE,
    COMMAND_CONFIGURE,
    COMMAND_IDENTITY,
    COMMAND_LIMITS,
    CONFIG_STRUCT,
    DEFAULT_REMOTE_HOST,
    DEFAULT_REMOTE_PORT,
    LIMITS_STRUCT,
    read_status,
    recv_exact,
    recv_float32_array,
    recv_u32,
)
from python_app.hardware_full.librevna_driver.models import SweepResult
from python_app.models.run_config_model import RadarSweepModel


@dataclass(slots=True)
class RemoteCompactMK209Service:
    """Acquire K209 sweeps through a persistent TCP connection."""

    host: str = DEFAULT_REMOTE_HOST
    port: int = DEFAULT_REMOTE_PORT
    timeout_s: float = 20.0
    _socket: socket.socket | None = field(init=False, default=None, repr=False)
    _settings: RadarSweepModel | None = field(init=False, default=None, repr=False)
    _frequency_hz: np.ndarray | None = field(init=False, default=None, repr=False)

    def __post_init__(self) -> None:
        self.host = str(self.host).strip()
        if not self.host:
            raise ValueError("K209 remote host must not be empty")
        self.port = int(self.port)
        if self.port <= 0 or self.port > 65535:
            raise ValueError("K209 remote port must be in 1..65535")
        self.timeout_s = float(self.timeout_s)
        if self.timeout_s <= 0.0:
            raise ValueError("K209 remote timeout_s must be > 0")

    def open(self) -> None:
        """Open TCP connection and apply stored settings when present."""
        if self._socket is not None:
            return
        sock = socket.create_connection((self.host, self.port), timeout=self.timeout_s)
        sock.settimeout(self.timeout_s)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self._socket = sock
        try:
            if self._settings is not None:
                self._apply_configuration(self._settings)
        except Exception:
            self.close()
            raise

    def close(self) -> None:
        """Close TCP connection."""
        if self._socket is None:
            return
        try:
            self._socket.close()
        finally:
            self._socket = None

    def query_identity(self) -> str:
        """Read analyzer identity string through the remote server."""
        sock = self._require_socket()
        sock.sendall(COMMAND_IDENTITY)
        read_status(sock)
        return recv_exact(sock, recv_u32(sock)).decode("utf-8", errors="replace").strip()

    def read_device_limits(self) -> dict[str, float | int]:
        """Read analyzer limits through the remote server."""
        opened_here = self._socket is None
        try:
            if opened_here:
                self.open()
            sock = self._require_socket()
            sock.sendall(COMMAND_LIMITS)
            read_status(sock)
            min_freq, max_freq, min_ifbw, max_ifbw, max_points, min_power, max_power = LIMITS_STRUCT.unpack(
                recv_exact(sock, LIMITS_STRUCT.size)
            )
            return {
                "min_frequency_hz": float(min_freq),
                "max_frequency_hz": float(max_freq),
                "min_ifbw_hz": float(min_ifbw),
                "max_ifbw_hz": float(max_ifbw),
                "max_points": int(max_points),
                "min_power_dbm": float(min_power),
                "max_power_dbm": float(max_power),
            }
        finally:
            if opened_here:
                self.close()

    def configure(self, sweep: RadarSweepModel) -> None:
        """Store and apply sweep settings."""
        self._validate_sweep(sweep)
        self._settings = sweep
        self._frequency_hz = None
        if self._socket is not None:
            self._apply_configuration(sweep)

    def acquire(self) -> SweepResult:
        """Acquire one corrected S11/S21 sweep."""
        if self._settings is None or self._frequency_hz is None:
            raise RuntimeError("K209 remote service is not configured")
        sock = self._require_socket()
        points = int(self._settings.points)
        sock.sendall(COMMAND_ACQUIRE)
        read_status(sock)
        returned_points = recv_u32(sock)
        if returned_points != points:
            raise RuntimeError(f"K209 remote sweep returned {returned_points} points, expected {points}")
        s11_values = recv_float32_array(sock, points * 2)
        s21_values = recv_float32_array(sock, points * 2)
        return SweepResult(
            x=self._frequency_hz.copy(),
            traces={
                "s11": self._complex_from_interleaved(s11_values),
                "s21": self._complex_from_interleaved(s21_values),
            },
        )

    def _apply_configuration(self, sweep: RadarSweepModel) -> None:
        sock = self._require_socket()
        sock.sendall(COMMAND_CONFIGURE)
        sock.sendall(
            CONFIG_STRUCT.pack(
                float(sweep.start_hz),
                float(sweep.stop_hz),
                int(sweep.points),
                float(sweep.if_bandwidth_hz),
                float(sweep.power_dbm),
            )
        )
        read_status(sock)
        returned_points = recv_u32(sock)
        if returned_points != int(sweep.points):
            raise RuntimeError(f"K209 remote config returned {returned_points} points, expected {sweep.points}")
        self._frequency_hz = recv_float32_array(sock, int(sweep.points))

    def _require_socket(self) -> socket.socket:
        if self._socket is None:
            raise RuntimeError("K209 remote socket is not open")
        return self._socket

    @staticmethod
    def _validate_sweep(sweep: RadarSweepModel) -> None:
        if int(sweep.points) < 2:
            raise ValueError("K209 sweep points must be >= 2")
        if float(sweep.stop_hz) < float(sweep.start_hz):
            raise ValueError("K209 sweep stop_hz must be >= start_hz")
        if float(sweep.if_bandwidth_hz) <= 0.0:
            raise ValueError("K209 IF bandwidth must be > 0")

    @staticmethod
    def _complex_from_interleaved(values: np.ndarray) -> np.ndarray:
        if values.size % 2 != 0:
            raise RuntimeError("K209 remote complex trace payload has odd scalar count")
        reshaped = np.asarray(values, dtype=np.float32).reshape((-1, 2))
        return (reshaped[:, 0] + 1j * reshaped[:, 1]).astype(np.complex64)
