radar_system/python_app/scripts/k209_remote_server.py

"""Run a TCP acquisition server for a locally connected Compact-M K209."""

from __future__ import annotations

import argparse
import socket
import socketserver
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from python_app.hardware_full.compact_m_k209_service import CompactMK209Service
from python_app.hardware_full.k209_remote_protocol import (
    COMMAND_ACQUIRE,
    COMMAND_CONFIGURE,
    COMMAND_IDENTITY,
    COMMAND_LIMITS,
    CONFIG_STRUCT,
    DEFAULT_REMOTE_PORT,
    LIMITS_STRUCT,
    STATUS_OK,
    recv_exact,
    send_error,
    send_float32_array,
    send_u32,
)
from python_app.models.run_config_model import RadarSweepModel

DEFAULT_RESOURCE = "TCPIP0::127.0.0.1::hislip0,4880::INSTR"


class K209RemoteRequestHandler(socketserver.StreamRequestHandler):
    """Handle one persistent K209 remote client connection."""

    def setup(self) -> None:
        """Disable Nagle and open the local K209 VISA session for this connection."""
        super().setup()
        self.request.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.service = CompactMK209Service(
            resource=self.server.resource,
            timeout_ms=self.server.timeout_ms,
            preset_on_open=False,
            visa_library="@ivi",
        )
        self.service.open()

    def finish(self) -> None:
        """Close the K209 VISA session when the client disconnects."""
        try:
            self.service.close()
        finally:
            super().finish()

    def handle(self) -> None:
        """Serve command bytes from the client until EOF, reporting errors back inline."""
        while True:
            command = self.rfile.read(1)
            if not command:
                return
            try:
                if command == COMMAND_IDENTITY:
                    self._handle_identity()
                elif command == COMMAND_LIMITS:
                    self._handle_limits()
                elif command == COMMAND_CONFIGURE:
                    self._handle_configure()
                elif command == COMMAND_ACQUIRE:
                    self._handle_acquire()
                else:
                    raise RuntimeError(f"unsupported command byte {command!r}")
                self.wfile.flush()
            except Exception as exc:  # noqa: BLE001
                send_error(self.wfile, str(exc))
                self.wfile.flush()

    def _handle_identity(self) -> None:
        """Reply with the device identity string (length-prefixed UTF-8)."""
        payload = self.service.query_identity().encode("utf-8")
        self.wfile.write(STATUS_OK)
        send_u32(self.wfile, len(payload))
        self.wfile.write(payload)

    def _handle_limits(self) -> None:
        """Reply with the device frequency/IFBW/power/point limits as a packed struct."""
        limits = self.service.read_device_limits()
        self.wfile.write(STATUS_OK)
        self.wfile.write(
            LIMITS_STRUCT.pack(
                float(limits["min_frequency_hz"]),
                float(limits["max_frequency_hz"]),
                float(limits["min_ifbw_hz"]),
                float(limits["max_ifbw_hz"]),
                int(limits["max_points"]),
                float(limits["min_power_dbm"]),
                float(limits["max_power_dbm"]),
            )
        )

    def _handle_configure(self) -> None:
        """Apply a sweep config from the client and reply with the frequency axis."""
        start_hz, stop_hz, points, ifbw_hz, power_dbm = CONFIG_STRUCT.unpack(
            recv_exact(self.rfile, CONFIG_STRUCT.size)
        )
        sweep = RadarSweepModel(
            start_hz=start_hz,
            stop_hz=stop_hz,
            points=int(points),
            if_bandwidth_hz=ifbw_hz,
            power_dbm=power_dbm,
        )
        self.service.configure(sweep)
        self.wfile.write(STATUS_OK)
        send_u32(self.wfile, int(points))
        send_float32_array(self.wfile, self.service.frequency_axis())

    def _handle_acquire(self) -> None:
        """Acquire one interleaved sweep and reply with the S11 and S21 arrays."""
        sweep = self.service.acquire_interleaved()
        self.wfile.write(STATUS_OK)
        send_u32(self.wfile, int(sweep.frequency_hz.size))
        send_float32_array(self.wfile, sweep.s11_values)
        send_float32_array(self.wfile, sweep.s21_values)


class K209RemoteServer(socketserver.TCPServer):
    """Single-client TCP server with K209 connection settings."""

    allow_reuse_address = True

    def __init__(self, server_address: tuple[str, int], resource: str, timeout_ms: int) -> None:
        """Store the VISA resource and timeout used by each accepted connection."""
        self.resource = resource
        self.timeout_ms = timeout_ms
        super().__init__(server_address, K209RemoteRequestHandler)


def _parse_args() -> argparse.Namespace:
    """Parse command-line options for the K209 remote server."""
    parser = argparse.ArgumentParser(description="Serve a locally connected Compact-M K209 over TCP.")
    parser.add_argument("--host", default="0.0.0.0", help="Server bind address.")
    parser.add_argument("--port", type=int, default=DEFAULT_REMOTE_PORT, help="Server TCP port.")
    parser.add_argument("--resource", default=DEFAULT_RESOURCE, help="Local S2VNA VISA resource.")
    parser.add_argument("--timeout-ms", type=int, default=20_000, help="K209 VISA timeout.")
    return parser.parse_args()


def main() -> int:
    """Bind the K209 remote server and serve clients until interrupted."""
    args = _parse_args()
    with K209RemoteServer((args.host, args.port), resource=args.resource, timeout_ms=args.timeout_ms) as server:
        print(f"K209 remote server listening on {args.host}:{args.port}")
        print(f"Local S2VNA resource: {args.resource}")
        server.serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())