ria-toolkit-oss/src/ria_toolkit_oss/agent/streamer.py

"""IQ-streaming agent.

Listens for control messages from the RIA Hub over a persistent WebSocket.
Supports:

- An **RX session** (hub sends ``start``/``stop``/``configure``; agent opens
  the SDR, loops ``sdr.rx()`` and ships raw interleaved float32 IQ).
- A **TX session** (hub sends ``tx_start``/``tx_stop``/``tx_configure`` plus
  binary IQ frames; agent feeds them into ``sdr._stream_tx``). Phase 3 wires
  up the session plumbing and rejects TX when ``cfg.tx_enabled`` is False;
  Phase 4 implements the full TX loop.

Both sessions can run concurrently on the same physical SDR (FDD) — a
ref-counted SDR registry shares one driver instance when RX and TX name the
same ``(device, identifier)``.
"""

from __future__ import annotations

import asyncio
import logging
import queue
import threading
import time
from dataclasses import dataclass, field
from typing import Any

import numpy as np

from .config import AgentConfig
from .hardware import heartbeat_payload
from .ws_client import WsClient

logger = logging.getLogger("ria_agent.streamer")

_DEFAULT_BUFFER_SIZE = 1024


# ---------------------------------------------------------------------------
# Session dataclasses


@dataclass
class RxSession:
    app_id: str
    sdr: Any
    device_key: tuple[str, str | None]
    buffer_size: int
    task: asyncio.Task | None = None
    pending_config: dict = field(default_factory=dict)


@dataclass
class TxSession:
    app_id: str
    sdr: Any
    device_key: tuple[str, str | None]
    buffer_size: int
    task: Any = None  # concurrent.futures.Future from run_in_executor
    pending_config: dict = field(default_factory=dict)
    underrun_policy: str = "pause"
    last_buffer: np.ndarray | None = None
    stop_event: threading.Event = field(default_factory=threading.Event)
    started_at: float = 0.0
    max_duration_s: float | None = None
    state: str = "armed"
    # Thread-safe queue of inbound interleaved-float32 IQ frames. Bounded so
    # hub-side over-production triggers WS backpressure rather than memory
    # growth in the agent.
    in_queue: "queue.Queue[bytes]" = field(default_factory=lambda: queue.Queue(maxsize=8))
    # Set by the TX callback when it hits an underrun while policy=="pause";
    # asyncio side flips the session state and emits tx_status.
    underrun_flag: threading.Event = field(default_factory=threading.Event)


# ---------------------------------------------------------------------------
# SDR registry (ref-counted so one Pluto handle serves RX + TX simultaneously)


class _SdrRegistry:
    def __init__(self, factory):
        self._factory = factory
        self._instances: dict[tuple[str, str | None], tuple[Any, int]] = {}
        self._lock = threading.Lock()

    def acquire(self, device: str, identifier: str | None) -> tuple[Any, tuple[str, str | None]]:
        key = (device, identifier)
        with self._lock:
            if key in self._instances:
                sdr, rc = self._instances[key]
                self._instances[key] = (sdr, rc + 1)
                return sdr, key
        # Build outside the lock: driver init can be slow and we don't want to
        # block concurrent releases on unrelated devices.
        sdr = self._factory(device, identifier)
        with self._lock:
            if key in self._instances:
                # Raced another acquirer; discard our duplicate and share theirs.
                other_sdr, rc = self._instances[key]
                try:
                    sdr.close()
                except Exception:
                    pass
                self._instances[key] = (other_sdr, rc + 1)
                return other_sdr, key
            self._instances[key] = (sdr, 1)
            return sdr, key

    def release(self, key: tuple[str, str | None]) -> bool:
        """Decrement refcount. Returns True if the caller owns the last reference
        and should close the SDR."""
        with self._lock:
            sdr, rc = self._instances.get(key, (None, 0))
            if sdr is None:
                return False
            if rc <= 1:
                del self._instances[key]
                return True
            self._instances[key] = (sdr, rc - 1)
            return False

    def refcount(self, key: tuple[str, str | None]) -> int:
        with self._lock:
            return self._instances.get(key, (None, 0))[1]


# ---------------------------------------------------------------------------
# Streamer


class Streamer:
    """Main streamer loop.

    Parameters
    ----------
    ws:
        Connected :class:`WsClient`.
    sdr_factory:
        Callable ``(device, identifier) -> SDR``. Defaults to the helper in
        :mod:`ria_toolkit_oss.sdr`. Injectable for tests.
    cfg:
        :class:`AgentConfig` for interlocks (``tx_enabled`` and caps) and
        heartbeat capabilities. Defaults to an empty ``AgentConfig()`` which
        leaves TX disabled.
    """

    def __init__(
        self,
        ws,
        sdr_factory=None,
        cfg: AgentConfig | None = None,
    ) -> None:
        self.ws = ws
        self._cfg = cfg or AgentConfig()
        self._registry = _SdrRegistry(sdr_factory or _default_sdr_factory)
        self._rx: RxSession | None = None
        self._tx: TxSession | None = None
        # Pending radio_config accepted via ``configure`` before ``start``.
        self._standalone_pending_config: dict = {}
        # Cached asyncio event loop, set the first time a handler runs. Used
        # to schedule async callbacks from the TX executor thread.
        self._loop: asyncio.AbstractEventLoop | None = None

    # ------------------------------------------------------------------
    # Back-compat read-only shims for callers that check ``._sdr`` etc.
    # Writes to these attributes are not supported — use the session objects.

    @property
    def _sdr(self):
        return self._rx.sdr if self._rx is not None else None

    @property
    def _pending_config(self) -> dict:
        return self._rx.pending_config if self._rx is not None else self._standalone_pending_config

    # ------------------------------------------------------------------
    # WsClient wiring

    def build_heartbeat(self) -> dict:
        status = "streaming" if (self._rx is not None or self._tx is not None) else "idle"
        app_id: str | None = None
        if self._rx is not None:
            app_id = self._rx.app_id
        elif self._tx is not None:
            app_id = self._tx.app_id

        sessions: dict[str, dict] = {}
        if self._rx is not None:
            sessions["rx"] = {"app_id": self._rx.app_id, "state": "streaming"}
        if self._tx is not None:
            sessions["tx"] = {"app_id": self._tx.app_id, "state": self._tx.state}

        return heartbeat_payload(
            status=status,
            app_id=app_id,
            cfg=self._cfg,
            sessions=sessions or None,
        )

    async def on_message(self, msg: dict) -> None:
        t = msg.get("type")
        handler = {
            "start": self._handle_rx_start,
            "stop": self._handle_rx_stop,
            "configure": self._handle_rx_configure,
            "tx_start": self._handle_tx_start,
            "tx_stop": self._handle_tx_stop,
            "tx_configure": self._handle_tx_configure,
        }.get(t)
        if handler is None:
            logger.warning("Unknown server message type: %r", t)
            return
        await handler(msg)

    async def on_binary(self, data: bytes) -> None:
        tx = self._tx
        if tx is None:
            logger.debug("Dropping %d-byte binary frame: no TX session", len(data))
            return
        # Backpressure: if the TX queue is full, await briefly so the hub's
        # ``await ws.send`` throttles naturally via TCP. We don't block
        # indefinitely — a 2s stall means something else is wrong.
        loop = asyncio.get_running_loop()
        try:
            await loop.run_in_executor(None, lambda: tx.in_queue.put(data, timeout=2.0))
        except queue.Full:
            logger.warning("TX queue stalled; dropping frame")

    # ==================================================================
    # RX

    async def _handle_rx_start(self, msg: dict) -> None:
        if self._rx is not None:
            logger.warning("start received while already streaming — ignoring")
            return

        app_id = msg.get("app_id") or ""
        radio_config = dict(msg.get("radio_config") or {})
        device = radio_config.pop("device", None)
        identifier = radio_config.pop("identifier", None)
        buffer_size = int(radio_config.pop("buffer_size", _DEFAULT_BUFFER_SIZE))
        if not device:
            await self._send_error(app_id, "start missing radio_config.device")
            return

        try:
            sdr, device_key = self._registry.acquire(device, identifier)
            _apply_sdr_config(sdr, radio_config)
        except Exception as exc:
            logger.exception("Failed to open SDR %r", device)
            await self._send_error(app_id, f"SDR init failed: {exc}")
            return

        # Inherit any pending config that was queued before start.
        pending = dict(self._standalone_pending_config)
        self._standalone_pending_config = {}

        session = RxSession(
            app_id=app_id,
            sdr=sdr,
            device_key=device_key,
            buffer_size=buffer_size,
            pending_config=pending,
        )
        self._rx = session
        await self._send_status("streaming", app_id)
        session.task = asyncio.create_task(
            self._capture_loop(session), name="ria-streamer-capture"
        )

    async def _handle_rx_stop(self, msg: dict) -> None:
        session = self._rx
        if session is None:
            return
        if session.task is not None:
            session.task.cancel()
            try:
                await session.task
            except (asyncio.CancelledError, Exception):
                pass
        self._close_session_sdr(session)
        app_id = session.app_id
        self._rx = None
        await self._send_status("idle", app_id)

    async def _handle_rx_configure(self, msg: dict) -> None:
        cfg = dict(msg.get("radio_config") or {})
        if self._rx is not None:
            self._rx.pending_config.update(cfg)
        else:
            self._standalone_pending_config.update(cfg)
        logger.debug("Queued configure: %s", cfg)

    async def _capture_loop(self, session: RxSession) -> None:
        loop = asyncio.get_running_loop()
        try:
            while True:
                if session.pending_config:
                    cfg = session.pending_config
                    session.pending_config = {}
                    try:
                        _apply_sdr_config(session.sdr, cfg)
                    except Exception as exc:
                        logger.warning("Applying configure failed: %s", exc)

                try:
                    samples = await loop.run_in_executor(
                        None, session.sdr.rx, session.buffer_size
                    )
                except Exception as exc:
                    from ria_toolkit_oss.sdr import SdrDisconnectedError

                    if isinstance(exc, SdrDisconnectedError):
                        logger.warning("SDR disconnected: %s", exc)
                        await self._send_error(session.app_id, f"SDR disconnected: {exc}")
                    else:
                        logger.exception("SDR rx error")
                        await self._send_error(session.app_id, f"SDR capture failed: {exc}")
                    break

                payload = _samples_to_interleaved_float32(samples)
                try:
                    await self.ws.send_bytes(payload)
                except Exception as exc:
                    logger.warning("Send failed: %s — ending capture", exc)
                    break
        except asyncio.CancelledError:
            raise
        finally:
            self._close_session_sdr(session)
            # If the loop died on its own (e.g. SDR disconnect), clear the
            # session handle so future ``start`` messages can proceed.
            if self._rx is session:
                self._rx = None

    # ==================================================================
    # TX

    async def _handle_tx_start(self, msg: dict) -> None:
        app_id = msg.get("app_id") or ""
        radio_config = dict(msg.get("radio_config") or {})

        # --- interlocks (agent-enforced; never trust the hub alone) ---
        if not self._cfg.tx_enabled:
            await self._send_tx_status(app_id, "error", "tx disabled on this agent")
            return
        tx_gain = radio_config.get("tx_gain")
        if (
            self._cfg.tx_max_gain_db is not None
            and tx_gain is not None
            and float(tx_gain) > float(self._cfg.tx_max_gain_db)
        ):
            await self._send_tx_status(
                app_id,
                "error",
                f"tx_gain {tx_gain} exceeds cap {self._cfg.tx_max_gain_db}",
            )
            return
        tx_freq = radio_config.get("tx_center_frequency")
        if self._cfg.tx_allowed_freq_ranges and tx_freq is not None:
            f = float(tx_freq)
            if not any(float(lo) <= f <= float(hi) for lo, hi in self._cfg.tx_allowed_freq_ranges):
                await self._send_tx_status(
                    app_id,
                    "error",
                    f"tx_center_frequency {tx_freq} outside allowed ranges",
                )
                return

        if self._tx is not None:
            await self._send_tx_status(app_id, "error", "tx already active on this agent")
            return

        # --- device ---
        device = radio_config.pop("device", None)
        identifier = radio_config.pop("identifier", None)
        buffer_size = int(radio_config.pop("buffer_size", _DEFAULT_BUFFER_SIZE))
        underrun_policy = str(radio_config.pop("underrun_policy", "pause"))
        if underrun_policy not in ("pause", "zero", "repeat"):
            await self._send_tx_status(
                app_id, "error", f"invalid underrun_policy {underrun_policy!r}"
            )
            return
        if not device:
            await self._send_tx_status(app_id, "error", "tx_start missing radio_config.device")
            return

        device_key: tuple[str, str | None] | None = None
        sdr: Any = None
        try:
            sdr, device_key = self._registry.acquire(device, identifier)
            _apply_sdr_config(sdr, radio_config)
            # Only call init_tx when the hub supplied the three required
            # parameters. Drivers that gate _stream_tx on _tx_initialized
            # (e.g. Pluto) need this; drivers that don't (e.g. Mock) tolerate
            # its absence.
            init_args = {
                k: radio_config.get(f"tx_{k}")
                for k in ("sample_rate", "center_frequency", "gain")
            }
            if hasattr(sdr, "init_tx") and all(v is not None for v in init_args.values()):
                sdr.init_tx(
                    sample_rate=init_args["sample_rate"],
                    center_frequency=init_args["center_frequency"],
                    gain=init_args["gain"],
                    channel=radio_config.get("tx_channel", 0),
                    gain_mode=radio_config.get("tx_gain_mode", "manual"),
                )
        except Exception as exc:
            if device_key is not None:
                if self._registry.release(device_key):
                    try:
                        sdr.close()
                    except Exception:
                        pass
            logger.exception("Failed to init TX on %r", device)
            await self._send_tx_status(app_id, "error", f"tx init failed: {exc}")
            return

        self._loop = asyncio.get_running_loop()
        session = TxSession(
            app_id=app_id,
            sdr=sdr,
            device_key=device_key,
            buffer_size=buffer_size,
            underrun_policy=underrun_policy,
            started_at=time.monotonic(),
            max_duration_s=self._cfg.tx_max_duration_s,
        )
        self._tx = session
        await self._send_tx_status(app_id, "armed")
        session.task = self._loop.run_in_executor(None, self._tx_executor_body, session)
        # Spawn a small watchdog that transitions armed → transmitting when
        # the first buffer has been consumed, and surfaces underrun / max-
        # duration terminations back to the hub.
        asyncio.create_task(self._tx_watchdog(session))

    async def _handle_tx_stop(self, msg: dict) -> None:
        session = self._tx
        if session is None:
            return
        app_id = session.app_id
        session.stop_event.set()
        try:
            session.sdr.pause_tx()
        except Exception:
            logger.debug("pause_tx raised during stop", exc_info=True)
        # Wake the executor thread if it's blocked on ``queue.get``.
        self._drain_tx_queue(session)
        if session.task is not None:
            try:
                await asyncio.wait_for(asyncio.wrap_future(session.task), timeout=1.5)
            except asyncio.TimeoutError:
                logger.warning("TX executor did not exit within 1.5s after stop")
            except Exception:
                logger.debug("TX executor raised on shutdown", exc_info=True)
        self._close_session_sdr(session)
        self._tx = None
        await self._send_tx_status(app_id, "done")

    async def _handle_tx_configure(self, msg: dict) -> None:
        if self._tx is None:
            return
        self._tx.pending_config.update(msg.get("radio_config") or {})

    # ------------------------------------------------------------------
    # TX executor & watchdog

    def _tx_executor_body(self, session: TxSession) -> None:
        try:
            session.sdr._stream_tx(lambda n: self._tx_callback(session, n))
        except Exception:
            logger.exception("TX stream crashed")
            self._schedule(self._send_tx_status(session.app_id, "error", "tx stream crashed"))

    def _tx_callback(self, session: TxSession, num_samples) -> np.ndarray:
        n = int(num_samples)
        # Honor stop requests: return silence one last time and let the driver
        # exit its loop on the next iteration (pause_tx flips _enable_tx).
        if session.stop_event.is_set():
            return _silence(n)

        # Max-duration watchdog.
        if (
            session.max_duration_s is not None
            and (time.monotonic() - session.started_at) >= float(session.max_duration_s)
        ):
            session.stop_event.set()
            try:
                session.sdr.pause_tx()
            except Exception:
                pass
            self._schedule(self._send_tx_status(session.app_id, "done", "max duration reached"))
            return _silence(n)

        # Apply queued configure at buffer boundary.
        if session.pending_config:
            cfg = session.pending_config
            session.pending_config = {}
            try:
                _apply_sdr_config(session.sdr, cfg)
            except Exception as exc:
                logger.debug("tx_configure apply failed: %s", exc)

        try:
            raw = session.in_queue.get(timeout=0.1)
        except queue.Empty:
            return self._underrun_fill(session, n)

        arr = np.frombuffer(raw, dtype=np.float32)
        if arr.size < 2 or arr.size % 2 != 0:
            logger.warning("Malformed TX frame: %d floats (must be non-zero even count)", arr.size)
            return self._underrun_fill(session, n)
        samples = (arr[0::2].astype(np.complex64) + 1j * arr[1::2].astype(np.complex64))
        if samples.size < n:
            out = np.zeros(n, dtype=np.complex64)
            out[: samples.size] = samples
            session.last_buffer = out
            return out
        if samples.size > n:
            samples = samples[:n]
        session.last_buffer = samples
        if session.state == "armed":
            session.state = "transmitting"
            self._schedule(self._send_tx_status(session.app_id, "transmitting"))
        return samples

    def _underrun_fill(self, session: TxSession, n: int) -> np.ndarray:
        policy = session.underrun_policy
        if policy == "zero":
            return _silence(n)
        if policy == "repeat" and session.last_buffer is not None:
            buf = session.last_buffer
            if buf.size == n:
                return buf
            if buf.size > n:
                return buf[:n].copy()
            out = np.zeros(n, dtype=np.complex64)
            out[: buf.size] = buf
            return out
        # "pause" policy (default) or "repeat" before any buffer arrived.
        if not session.underrun_flag.is_set():
            session.underrun_flag.set()
        session.stop_event.set()
        try:
            session.sdr.pause_tx()
        except Exception:
            pass
        return _silence(n)

    async def _tx_watchdog(self, session: TxSession) -> None:
        # Poll the underrun flag so we can emit status + tear down cleanly
        # when the callback flips the flag from the executor thread. Check
        # underrun_flag before stop_event, since the "pause" path sets both.
        while session is self._tx:
            if session.underrun_flag.is_set():
                await self._send_tx_status(session.app_id, "underrun")
                await self._teardown_tx_after_underrun(session)
                return
            if session.stop_event.is_set():
                return
            await asyncio.sleep(0.05)

    async def _teardown_tx_after_underrun(self, session: TxSession) -> None:
        if self._tx is not session:
            return
        self._drain_tx_queue(session)
        if session.task is not None:
            try:
                await asyncio.wait_for(asyncio.wrap_future(session.task), timeout=1.0)
            except asyncio.TimeoutError:
                logger.warning("TX executor did not exit within 1s after underrun")
            except Exception:
                logger.debug("TX executor raised during underrun teardown", exc_info=True)
        self._close_session_sdr(session)
        if self._tx is session:
            self._tx = None

    def _drain_tx_queue(self, session: TxSession) -> None:
        try:
            while True:
                session.in_queue.get_nowait()
        except queue.Empty:
            pass

    def _schedule(self, coro) -> None:
        loop = self._loop
        if loop is None:
            return
        try:
            asyncio.run_coroutine_threadsafe(coro, loop)
        except Exception:
            logger.debug("_schedule failed", exc_info=True)

    # ==================================================================
    # Helpers

    def _close_session_sdr(self, session) -> None:
        if session.sdr is None:
            return
        should_close = self._registry.release(session.device_key)
        if should_close:
            try:
                session.sdr.close()
            except Exception:
                logger.debug("SDR close raised", exc_info=True)

    async def _send_status(self, status: str, app_id: str) -> None:
        try:
            await self.ws.send_json({"type": "status", "status": status, "app_id": app_id})
        except Exception as exc:
            logger.debug("Status send failed: %s", exc)

    async def _send_error(self, app_id: str, message: str) -> None:
        try:
            await self.ws.send_json({"type": "error", "app_id": app_id, "message": message})
        except Exception as exc:
            logger.debug("Error-frame send failed: %s", exc)

    async def _send_tx_status(self, app_id: str, state: str, message: str | None = None) -> None:
        payload: dict = {"type": "tx_status", "app_id": app_id, "state": state}
        if message is not None:
            payload["message"] = message
        try:
            await self.ws.send_json(payload)
        except Exception as exc:
            logger.debug("tx_status send failed: %s", exc)


# ---------------------------------------------------------------------------
# Helpers

_CONFIG_ATTR_MAP = {
    "sample_rate": ("sample_rate", "rx_sample_rate"),
    "center_frequency": ("center_freq", "rx_center_frequency"),
    "center_freq": ("center_freq", "rx_center_frequency"),
    "gain": ("gain", "rx_gain"),
    "bandwidth": ("bandwidth", "rx_bandwidth"),
    "tx_sample_rate": ("tx_sample_rate",),
    "tx_center_frequency": ("tx_center_frequency", "tx_lo"),
    "tx_gain": ("tx_gain",),
    "tx_bandwidth": ("tx_bandwidth",),
}


def _apply_sdr_config(sdr: Any, cfg: dict) -> None:
    """Apply a radio_config dict to an SDR, trying multiple attribute aliases."""
    for key, value in cfg.items():
        if value is None:
            continue
        attrs = _CONFIG_ATTR_MAP.get(key, (key,))
        applied = False
        for attr in attrs:
            if hasattr(sdr, attr):
                try:
                    setattr(sdr, attr, value)
                    applied = True
                    break
                except Exception as exc:
                    logger.debug("setattr %s=%r failed: %s", attr, value, exc)
        if not applied:
            logger.debug("radio_config key %r ignored (no matching attr)", key)


def _silence(num_samples: int) -> np.ndarray:
    """Return a ``num_samples``-length zero-filled complex64 buffer."""
    return np.zeros(int(num_samples), dtype=np.complex64)


def _samples_to_interleaved_float32(samples: Any) -> bytes:
    """Convert complex IQ samples (any numeric dtype) to interleaved float32 bytes."""
    arr = np.asarray(samples)
    if np.iscomplexobj(arr):
        interleaved = np.empty(arr.size * 2, dtype=np.float32)
        interleaved[0::2] = arr.real.astype(np.float32, copy=False).ravel()
        interleaved[1::2] = arr.imag.astype(np.float32, copy=False).ravel()
        return interleaved.tobytes()
    return arr.astype(np.float32, copy=False).tobytes()


def _default_sdr_factory(device: str, identifier: str | None):
    from ria_toolkit_oss.sdr import get_sdr_device

    return get_sdr_device(device, ident=identifier)


# ---------------------------------------------------------------------------
# Top-level entry

async def run_streamer(ws_url: str, token: str, *, cfg: AgentConfig | None = None) -> None:
    """Connect to *ws_url* and run the streamer loop until cancelled."""
    ws = WsClient(ws_url, token)
    streamer = Streamer(ws, cfg=cfg)
    await ws.run(
        streamer.on_message,
        streamer.build_heartbeat,
        on_binary=streamer.on_binary,
    )