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

"""Unified ``ria-agent`` CLI.

Subcommands:

- ``ria-agent run [legacy args]`` — legacy long-poll NodeAgent (unchanged).
- ``ria-agent stream`` — new WebSocket-based IQ streamer.
- ``ria-agent detect`` — print SDR drivers whose modules import cleanly.
- ``ria-agent register --hub URL --api-key KEY`` — register with the hub and
  save credentials (and optional TX interlocks) to ``~/.ria/agent.json``.

Invoking ``ria-agent`` with no subcommand falls through to the legacy
long-poll behavior for back-compatibility with existing deployments.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import sys

from . import config as _config
from .hardware import available_devices
from .legacy_executor import main as _legacy_main

_LEGACY_ALIASES = {"--hub", "--key", "--name", "--device", "--insecure", "--log-level", "--config"}


def _cmd_detect(_args: argparse.Namespace) -> int:
    devices = available_devices()
    if not devices:
        print("No SDR drivers available (install ria-toolkit-oss[all-sdr] or per-driver extras).")
        return 0
    for name in devices:
        print(name)
    return 0


def _cmd_register(args: argparse.Namespace) -> int:
    import urllib.request

    hub_url = args.hub.rstrip("/")
    url = f"{hub_url}/screens/agents/register"
    body = json.dumps({"name": args.name or ""}).encode()
    req = urllib.request.Request(
        url,
        data=body,
        headers={
            "Content-Type": "application/json",
            "X-API-Key": args.api_key,
        },
    )
    try:
        with urllib.request.urlopen(req) as resp:
            data = json.loads(resp.read())
    except Exception as e:
        print(f"error: registration failed: {e}", file=sys.stderr)
        return 1

    agent_id = data["agent_id"]
    token = data["token"]

    cfg = _config.load()
    cfg.hub_url = hub_url
    cfg.agent_id = agent_id
    cfg.token = token
    cfg.api_key = args.api_key
    if args.name:
        cfg.name = args.name
    cfg.insecure = bool(args.insecure)
    cfg.tx_enabled = bool(getattr(args, "allow_tx", False))
    if (v := getattr(args, "tx_max_gain_db", None)) is not None:
        cfg.tx_max_gain_db = float(v)
    if (v := getattr(args, "tx_max_duration_s", None)) is not None:
        cfg.tx_max_duration_s = float(v)
    freq_ranges = getattr(args, "tx_freq_range", None) or []
    if freq_ranges:
        cfg.tx_allowed_freq_ranges = [[float(lo), float(hi)] for lo, hi in freq_ranges]
    path = _config.save(cfg)

    print(f"Registered agent: {agent_id}")
    if cfg.tx_enabled:
        caps: list[str] = []
        if cfg.tx_max_gain_db is not None:
            caps.append(f"gain<={cfg.tx_max_gain_db} dB")
        if cfg.tx_max_duration_s is not None:
            caps.append(f"duration<={cfg.tx_max_duration_s} s")
        if cfg.tx_allowed_freq_ranges:
            caps.append(f"freq in {cfg.tx_allowed_freq_ranges}")
        tail = f" ({', '.join(caps)})" if caps else ""
        print(f"TX enabled{tail}")
    print(f"Credentials saved to {path}")
    return 0


def _cmd_stream(args: argparse.Namespace) -> int:
    from .streamer import run_streamer

    cfg = _config.load()
    url = args.url or _derive_ws_url(cfg.hub_url, cfg.agent_id)
    token = args.token or cfg.token
    if not url:
        print("error: --url is required (or run `ria-agent register` first)", file=sys.stderr)
        return 2
    if getattr(args, "allow_tx", False):
        cfg.tx_enabled = True
    try:
        asyncio.run(run_streamer(url, token, cfg=cfg))
    except KeyboardInterrupt:
        pass
    return 0


def _derive_ws_url(hub_url: str, agent_id: str) -> str:
    if not hub_url:
        return ""
    base = hub_url.rstrip("/")
    if base.startswith("https://"):
        base = "wss://" + base[len("https://"):]
    elif base.startswith("http://"):
        base = "ws://" + base[len("http://"):]
    suffix = f"/screens/agent/ws?agent_id={agent_id}" if agent_id else "/screens/agent/ws"
    return base + suffix


def main() -> None:
    # Back-compat: if the first non-flag token matches a known legacy flag,
    # or there is no subcommand at all, dispatch to the legacy CLI.
    argv = sys.argv[1:]
    if not argv or (argv[0].startswith("--") and argv[0] in _LEGACY_ALIASES):
        _legacy_main()
        return

    parser = argparse.ArgumentParser(prog="ria-agent")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("run", help="Legacy long-poll agent (NodeAgent)")
    sub.add_parser("detect", help="List available SDR drivers")

    p_reg = sub.add_parser("register", help="Register agent with RIA Hub and save credentials")
    p_reg.add_argument("--hub", required=True, help="RIA Hub URL (e.g. http://whitehorse:3005)")
    p_reg.add_argument("--api-key", dest="api_key", required=True, help="Hub API key")
    p_reg.add_argument("--name", default=None, help="Human-friendly agent name")
    p_reg.add_argument("--insecure", action="store_true", help="Skip TLS verification")
    p_reg.add_argument(
        "--allow-tx",
        dest="allow_tx",
        action="store_true",
        help="Opt this agent in to TX (required for any transmission from the hub)",
    )
    p_reg.add_argument(
        "--tx-max-gain-db",
        dest="tx_max_gain_db",
        type=float,
        default=None,
        help="Reject tx_start frames whose tx_gain exceeds this cap (dB)",
    )
    p_reg.add_argument(
        "--tx-max-duration-s",
        dest="tx_max_duration_s",
        type=float,
        default=None,
        help="Auto-stop any TX session after this many seconds",
    )
    p_reg.add_argument(
        "--tx-freq-range",
        dest="tx_freq_range",
        type=float,
        nargs=2,
        action="append",
        metavar=("LO", "HI"),
        default=None,
        help="Allowed TX center-frequency range in Hz (repeat for multiple bands)",
    )

    p_stream = sub.add_parser("stream", help="Run the WebSocket IQ streamer")
    p_stream.add_argument("--url", default=None, help="Override WebSocket URL")
    p_stream.add_argument("--token", default=None, help="Override bearer token")
    p_stream.add_argument("--log-level", default="INFO")
    p_stream.add_argument(
        "--allow-tx",
        dest="allow_tx",
        action="store_true",
        help="Runtime override: enable TX for this process without writing config",
    )

    # Unknown extras are forwarded to the legacy CLI when command == "run".
    args, extras = parser.parse_known_args(argv)

    logging.basicConfig(
        level=getattr(logging, getattr(args, "log_level", "INFO"), logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    if args.command == "run":
        sys.argv = [sys.argv[0], *extras]
        _legacy_main()
        return
    if args.command == "detect":
        sys.exit(_cmd_detect(args))
    if args.command == "register":
        sys.exit(_cmd_register(args))
    if args.command == "stream":
        sys.exit(_cmd_stream(args))

    parser.error(f"unknown command: {args.command}")


if __name__ == "__main__":
    main()