bandsaunter/bandsaunter/device.py

"""High level RTL-SDR device control built on the ctypes binding."""

from __future__ import annotations

import contextlib
import ctypes
import math
import os
import threading
import time
from collections import deque
from dataclasses import dataclass, field

import numpy as np

from ._quiet import suppress_stderr
from .librtlsdr import READ_ASYNC_CB, LibRtlSdr, RtlSdrError

__all__ = ["RtlSdrDevice", "DeviceInfo", "list_devices", "RtlSdrError",
           "quiet_driver", "set_driver_messages",
           "R820T_MIN_HZ", "R820T_MAX_HZ", "DIRECT_SAMPLING_MAX_HZ"]

# Practical tuning limits for the common R820T/R828D front ends.  Anything
# below this needs direct sampling (bypassing the tuner entirely).
R820T_MIN_HZ = 24_000_000
R820T_MAX_HZ = 1_766_000_000
DIRECT_SAMPLING_MAX_HZ = 28_800_000  # usable Q-branch span at 2.4 MSPS-ish clock

# The RTL2832U only locks these sample-rate windows.
_RATE_WINDOWS = ((225_001, 300_000), (900_001, 3_200_000))

# librtlsdr wants read sizes that are a multiple of 512 bytes; 16 kB blocks are
# the size its own async reader uses and behave well over USB.
_READ_GRANULE = 512


# librtlsdr writes its own messages ("Found Rafael Micro R820T tuner",
# "Allocating 15 zero-copy buffers", "[R82XX] PLL not locked!") straight to
# file descriptor 2 from C.  Those land in the middle of the live display and
# break its cursor tracking, so the header ends up drawn several times over.
# Set BANDSAUNTER_DRIVER_MESSAGES=1 to see them while debugging.
_SHOW_DRIVER_MESSAGES = os.environ.get("BANDSAUNTER_DRIVER_MESSAGES", "") not in \
    ("", "0", "no", "false")


def set_driver_messages(show: bool) -> None:
    """Let the driver's own messages through again.

    They are worth seeing while diagnosing a device -- "Detached kernel
    driver", the tuner it found -- and only a nuisance during the live
    display, which they draw straight over.
    """
    global _SHOW_DRIVER_MESSAGES
    _SHOW_DRIVER_MESSAGES = bool(show)


@contextlib.contextmanager
def quiet_driver():
    """Silence the driver's own chatter for the duration of a C call.

    Only around the call itself, so a Python traceback on stderr is never
    swallowed: the window is a few microseconds of library code.
    """
    with suppress_stderr(not _SHOW_DRIVER_MESSAGES):
        yield


@dataclass
class DeviceInfo:
    index: int
    name: str
    manufacturer: str = "?"
    product: str = "?"
    serial: str = "?"

    def describe(self) -> str:
        return (f"[{self.index}] {self.name} "
                f"(sn: {self.serial})")


def list_devices() -> list[DeviceInfo]:
    """Enumerate attached RTL-SDR devices (empty list if none / no driver)."""
    try:
        lib = LibRtlSdr()
    except RtlSdrError:
        return []
    out = []
    with quiet_driver():
        for i in range(lib.device_count()):
            manu, prod, serial = lib.usb_strings(i)
            out.append(DeviceInfo(i, lib.device_name(i), manu, prod, serial))
    return out


def clamp_sample_rate(rate: int) -> int:
    """Snap a requested sample rate into a window the RTL2832U can actually lock."""
    rate = int(rate)
    for lo, hi in _RATE_WINDOWS:
        if lo <= rate <= hi:
            return rate
    # Choose the nearest legal edge.
    edges = [e for w in _RATE_WINDOWS for e in w]
    return min(edges, key=lambda e: abs(e - rate))


@dataclass
class RtlSdrDevice:
    """Owns one dongle and serialises all access to it.

    The scanner retunes constantly, so this class keeps track of the settling
    cost of each retune and flushes the USB pipeline so that samples handed
    back are guaranteed to have been captured *after* the tune completed.
    """

    index: int = 0
    sample_rate: int = 2_048_000
    gain: float | str = "auto"          # dB, or "auto" for the tuner AGC
    ppm: int = 0
    agc: bool = False                    # RTL2832 digital AGC
    bias_tee: bool = False
    offset_tuning: bool = False
    direct_sampling: int | str = "auto"  # 0/1/2 or "auto" (enable below 24 MHz)
    settle_seconds: float = 0.006

    _lib: LibRtlSdr | None = field(default=None, init=False, repr=False)
    _dev: ctypes.c_void_p | None = field(default=None, init=False, repr=False)
    _buf: ctypes.Array | None = field(default=None, init=False, repr=False)
    _buf_len: int = field(default=0, init=False, repr=False)
    _ds_mode: int = field(default=0, init=False, repr=False)
    _center: int = field(default=0, init=False, repr=False)
    _gains: list[float] = field(default_factory=list, init=False, repr=False)
    tuner: str = field(default="unknown", init=False)
    retunes: int = field(default=0, init=False)
    samples_read: int = field(default=0, init=False)
    _streaming: bool = field(default=False, init=False, repr=False)
    _stream_thread: object = field(default=None, init=False, repr=False)
    _stream_cb: object = field(default=None, init=False, repr=False)
    _stream_chunks: object = field(default=None, init=False, repr=False)
    _stream_lock: object = field(default=None, init=False, repr=False)
    _stream_ready: object = field(default=None, init=False, repr=False)
    _stream_error: object = field(default=None, init=False, repr=False)
    _stream_have: int = field(default=0, init=False, repr=False)
    _stream_max: int = field(default=0, init=False, repr=False)
    _stream_dropped: int = field(default=0, init=False, repr=False)
    _partial: object = field(default=None, init=False, repr=False)

    # -- lifecycle -----------------------------------------------------
    def open(self) -> "RtlSdrDevice":
        self._lib = LibRtlSdr()
        n = self._lib.device_count()
        if n == 0:
            raise RtlSdrError(
                "no RTL-SDR device found.\n"
                "Check `lsusb` for a Realtek RTL2832/2838, make sure the DVB-T "
                "kernel module is blacklisted, and that you have permission to "
                "open the USB device."
            )
        if self.index >= n:
            raise RtlSdrError(f"device index {self.index} out of range (found {n})")

        with quiet_driver():
            self._dev = self._lib.open(self.index)
            self.tuner = self._lib.get_tuner_type(self._dev)
            self._gains = self._lib.get_tuner_gains(self._dev)

            self.sample_rate = clamp_sample_rate(self.sample_rate)
            self._lib.set_sample_rate(self._dev, self.sample_rate)
            self.sample_rate = (self._lib.get_sample_rate(self._dev)
                                or self.sample_rate)

            if self.ppm:
                self._lib.set_freq_correction(self._dev, self.ppm)
            self._lib.set_agc_mode(self._dev, self.agc)
            if self.bias_tee:
                self._lib.set_bias_tee(self._dev, True)
            if self.offset_tuning:
                self._lib.set_offset_tuning(self._dev, True)

            self.apply_gain(self.gain)
            self._lib.reset_buffer(self._dev)
        return self

    def close(self) -> None:
        self.stop_stream()
        if self._dev is not None and self._lib is not None:
            try:
                if self.bias_tee:
                    self._lib.set_bias_tee(self._dev, False)
            except RtlSdrError:
                pass
            with quiet_driver():
                self._lib.close(self._dev)
        self._dev = None

    def __enter__(self):
        return self.open()

    def __exit__(self, *exc):
        self.close()
        return False

    @property
    def is_open(self) -> bool:
        return self._dev is not None

    # -- configuration -------------------------------------------------
    @property
    def available_gains(self) -> list[float]:
        return list(self._gains)

    def apply_gain(self, gain: float | str) -> None:
        """Set tuner gain.  ``"auto"`` hands control to the tuner AGC."""
        self.gain = gain
        if isinstance(gain, str) and gain.lower() in ("auto", "agc", ""):
            self._lib.set_tuner_gain_mode(self._dev, False)
            return
        self._lib.set_tuner_gain_mode(self._dev, True)
        target = float(gain)
        if self._gains:
            target = min(self._gains, key=lambda g: abs(g - target))
        self._lib.set_tuner_gain(self._dev, target)

    @property
    def current_gain(self) -> float | str:
        if isinstance(self.gain, str):
            return "auto"
        try:
            return self._lib.get_tuner_gain(self._dev)
        except RtlSdrError:
            return self.gain

    def set_sample_rate(self, rate: int) -> int:
        rate = clamp_sample_rate(rate)
        if rate != self.sample_rate:
            with quiet_driver():
                self._lib.set_sample_rate(self._dev, rate)
                self.sample_rate = self._lib.get_sample_rate(self._dev) or rate
                self._lib.reset_buffer(self._dev)
        return self.sample_rate

    def _wants_direct_sampling(self, hz: float) -> int:
        if self.direct_sampling == "auto":
            return 2 if hz < R820T_MIN_HZ else 0
        return int(self.direct_sampling)

    def _set_direct_sampling(self, mode: int) -> None:
        if mode == self._ds_mode:
            return
        with quiet_driver():
            ok = self._lib.set_direct_sampling(self._dev, mode)
        if ok:
            self._ds_mode = mode
            # Switching the signal path invalidates whatever is in flight.
            self._lib.reset_buffer(self._dev)
        elif mode:
            raise RtlSdrError(
                "this dongle does not support direct sampling, so frequencies "
                f"below {R820T_MIN_HZ/1e6:.0f} MHz cannot be tuned"
            )

    def tune(self, hz: float, settle: bool = True) -> int:
        """Tune the front end.  Returns the frequency the hardware reports."""
        hz = int(round(hz))
        mode = self._wants_direct_sampling(hz)
        self._set_direct_sampling(mode)

        if mode == 0 and not (R820T_MIN_HZ <= hz <= R820T_MAX_HZ):
            raise RtlSdrError(
                f"{hz/1e6:.4f} MHz is outside the tuner's range "
                f"({R820T_MIN_HZ/1e6:.0f}-{R820T_MAX_HZ/1e6:.0f} MHz)"
            )

        with quiet_driver():
            self._lib.set_center_freq(self._dev, hz)
            self._center = self._lib.get_center_freq(self._dev) or hz
        self.retunes += 1
        if settle and self.settle_seconds > 0:
            time.sleep(self.settle_seconds)
        return self._center

    @property
    def center_freq(self) -> int:
        return self._center

    @property
    def direct_sampling_mode(self) -> int:
        """0 = normal tuner path, 1 = I branch, 2 = Q branch."""
        return self._ds_mode

    def can_tune(self, hz: float) -> bool:
        mode = self._wants_direct_sampling(hz)
        if mode:
            return 0 < hz <= DIRECT_SAMPLING_MAX_HZ
        return R820T_MIN_HZ <= hz <= R820T_MAX_HZ

    # -- capture -------------------------------------------------------
    def _ensure_buffer(self, n_bytes: int) -> ctypes.Array:
        if self._buf is None or self._buf_len < n_bytes:
            self._buf = (ctypes.c_ubyte * n_bytes)()
            self._buf_len = n_bytes
        return self._buf

    def flush(self) -> None:
        """Drop anything the RTL2832 already queued (call right after a retune)."""
        self._lib.reset_buffer(self._dev)

    def read_samples(self, count: int, flush: bool = False) -> np.ndarray:
        """Read ``count`` complex samples, returned as complex64 in [-1, 1).

        ``flush=True`` discards the in-flight USB buffer first so the samples
        are known to post-date the most recent tune.
        """
        if self._dev is None:
            raise RtlSdrError("device is not open")
        if flush:
            self._lib.reset_buffer(self._dev)

        n_bytes = int(count) * 2
        # Round up to the USB granule; librtlsdr short-reads otherwise.
        n_bytes = int(math.ceil(n_bytes / _READ_GRANULE) * _READ_GRANULE)
        buf = self._ensure_buffer(n_bytes)

        got = self._lib.read_sync(self._dev, buf, n_bytes)
        if got <= 0:
            raise RtlSdrError("USB read returned no data")

        raw = np.frombuffer(buf, dtype=np.uint8, count=got)
        if raw.size % 2:
            raw = raw[:-1]
        # 127.4 rather than 127.5 matches the RTL2832's actual DC bias.
        iq = raw.astype(np.float32)
        iq -= 127.4
        iq *= np.float32(1.0 / 128.0)
        out = iq.view(np.complex64) if iq.flags.c_contiguous else iq.copy().view(np.complex64)
        out = out[:count]
        self.samples_read += out.size
        return out

    def read_seconds(self, seconds: float, flush: bool = False) -> np.ndarray:
        return self.read_samples(int(self.sample_rate * seconds), flush=flush)

    # -- continuous streaming --------------------------------------------
    # Synchronous reads lose whatever the dongle sends between calls: the USB
    # host only moves data while a transfer is outstanding, so every
    # microsecond spent demodulating is a microsecond of samples thrown away.
    # The recording then holds less than it should and plays back too fast.
    # Streaming keeps a ring of transfers queued so the capture is gapless.

    def start_stream(self, buf_num: int = 15, buf_len: int = 262144,
                     max_seconds: float = 4.0) -> None:
        if self._streaming:
            return
        if self._dev is None:
            raise RtlSdrError("device is not open")
        self._stream_chunks = deque()
        self._stream_have = 0
        self._stream_max = int(max_seconds * self.sample_rate * 2)
        self._stream_dropped = 0
        self._stream_lock = threading.Lock()
        self._stream_ready = threading.Event()
        self._stream_error = None
        self._partial = np.zeros(0, dtype=np.complex64)

        def _on_buffer(buf_ptr, length, _ctx):
            try:
                raw = np.ctypeslib.as_array(buf_ptr, shape=(int(length),)).copy()
            except Exception:
                return
            with self._stream_lock:
                if self._stream_have + raw.size > self._stream_max:
                    # The consumer has fallen behind for real; drop the oldest
                    # so the capture stays current, and record that it happened.
                    while self._stream_chunks and \
                            self._stream_have + raw.size > self._stream_max:
                        old = self._stream_chunks.popleft()
                        self._stream_have -= old.size
                        self._stream_dropped += old.size
                self._stream_chunks.append(raw)
                self._stream_have += raw.size
            self._stream_ready.set()

        self._stream_cb = READ_ASYNC_CB(_on_buffer)

        def _run():
            try:
                with quiet_driver():
                    self._lib.read_async(self._dev, self._stream_cb,
                                         buf_num, buf_len)
            except RtlSdrError as exc:
                self._stream_error = exc
                self._stream_ready.set()

        with quiet_driver():
            self._lib.reset_buffer(self._dev)
        self._streaming = True
        self._stream_thread = threading.Thread(target=_run, daemon=True,
                                               name="rtlsdr-stream")
        self._stream_thread.start()

    def stop_stream(self) -> None:
        if not self._streaming:
            return
        self._streaming = False
        try:
            with quiet_driver():
                self._lib.cancel_async(self._dev)
        except Exception:
            pass
        if self._stream_thread is not None:
            self._stream_thread.join(timeout=2.0)
        self._stream_thread = None
        self._stream_cb = None

    @property
    def streaming(self) -> bool:
        return self._streaming

    @property
    def dropped_samples(self) -> int:
        return self._stream_dropped // 2

    def read_stream(self, count: int, timeout: float = 3.0) -> np.ndarray:
        """Pull ``count`` complex samples from the running stream."""
        if not self._streaming:
            return self.read_samples(count)
        need_bytes = int(count) * 2
        deadline = time.time() + timeout
        parts = []
        have = 0
        if self._partial.size:
            parts.append(self._partial)
            have += self._partial.size * 2
            self._partial = np.zeros(0, dtype=np.complex64)

        while have < need_bytes:
            if self._stream_error is not None:
                raise self._stream_error
            with self._stream_lock:
                chunk = self._stream_chunks.popleft() if self._stream_chunks else None
                if chunk is not None:
                    self._stream_have -= chunk.size
                elif not self._stream_chunks:
                    self._stream_ready.clear()
            if chunk is None:
                if time.time() > deadline:
                    raise RtlSdrError("timed out waiting for samples")
                self._stream_ready.wait(0.2)
                continue
            parts.append(self._to_complex(chunk))
            have += chunk.size

        out = np.concatenate([p if p.dtype == np.complex64 else p for p in parts])
        if out.size > count:
            self._partial = out[count:]
            out = out[:count]
        self.samples_read += out.size
        return out

    @staticmethod
    def _to_complex(raw: np.ndarray) -> np.ndarray:
        if raw.size % 2:
            raw = raw[:-1]
        iq = raw.astype(np.float32)
        iq -= 127.4
        iq *= np.float32(1.0 / 128.0)
        return iq.view(np.complex64)

    # -- diagnostics ---------------------------------------------------
    def status(self) -> dict:
        return {
            "index": self.index,
            "tuner": self.tuner,
            "sample_rate": self.sample_rate,
            "center_freq": self._center,
            "gain": self.current_gain,
            "ppm": self.ppm,
            "direct_sampling": self._ds_mode,
            "retunes": self.retunes,
            "samples_read": self.samples_read,
        }