ria-toolkit-oss/src/ria_toolkit_oss_cli/ria_toolkit_oss/transmit.py

"""Transmit command for SDR devices."""

import os
import signal
import time

import click

from ria_toolkit_oss.datatypes import Recording
from ria_toolkit_oss.io import from_npy_legacy, load_recording

from .common import (
    echo_progress,
    echo_verbose,
    format_frequency,
    format_sample_rate,
    get_sdr_device,
    load_yaml_config,
    parse_frequency,
)
from .discover import (
    find_bladerf_devices,
    find_hackrf_devices,
    find_pluto_devices,
    find_uhd_devices,
    load_sdr_drivers,
)

# TX-capable devices (RTL-SDR and ThinkRF are RX-only)
TX_CAPABLE_DEVICES = ["pluto", "hackrf", "bladerf", "usrp"]


def auto_select_tx_device(quiet: bool = False) -> str:
    """
    Auto-select TX-capable device if only one is connected.

    Args:
        quiet: Suppress warning messages

    Returns:
        Device type string

    Raises:
        click.ClickException: If no TX devices or multiple devices found
    """
    # Load drivers and collect TX-capable devices only
    load_sdr_drivers(verbose=False)

    tx_devices = []
    tx_devices.extend(find_uhd_devices())
    tx_devices.extend(find_pluto_devices())
    tx_devices.extend(find_hackrf_devices())
    tx_devices.extend(find_bladerf_devices())
    # Note: RTL-SDR and ThinkRF excluded (RX-only)

    if len(tx_devices) == 0:
        raise click.ClickException(
            "No TX-capable SDR devices found.\n"
            "TX-capable devices: PlutoSDR, HackRF, BladeRF, USRP\n"
            "Run 'ria discover' to see all devices."
        )

    elif len(tx_devices) == 1:
        device = tx_devices[0]
        device_type = device.get("type", "Unknown").lower().replace("-", "").replace(" ", "")

        # Map device type names to internal names
        type_map = {
            "plutosdr": "pluto",
            "hackrf": "hackrf",
            "hackrfone": "hackrf",
            "bladerf": "bladerf",
            "usrp": "usrp",
            "b200": "usrp",
            "b210": "usrp",
        }

        device_type = type_map.get(device_type, device_type)

        if not quiet:
            click.echo(
                click.style("Warning: ", fg="yellow")
                + f"No device specified. Auto-detected {device.get('type', 'Unknown')}",
                err=True,
            )
            click.echo(f"Use --device {device_type} to suppress this warning.\n", err=True)

        return device_type

    else:
        device_list = "\n".join(f"  - {d.get('type', 'Unknown')}" for d in tx_devices)
        raise click.ClickException(
            f"Multiple TX-capable devices found. Specify with --device\n\n"
            f"Available TX devices:\n{device_list}\n\n"
            f"Run 'ria discover' for more details."
        )


def load_input_file(input_file: str, legacy: bool = False) -> Recording:
    """
    Load recording from file with auto-format detection.

    Args:
        input_file: Path to input file
        legacy: Use legacy NPY loader

    Returns:
        Recording object

    Raises:
        click.ClickException: If file not found or format unsupported
    """
    if not os.path.exists(input_file):
        raise click.ClickException(f"Input file not found: {input_file}")

    try:
        if legacy:
            echo_progress("Loading legacy NPY file...", quiet=False)
            recording = from_npy_legacy(input_file)
        else:
            echo_progress("Loading input file...", quiet=False)
            recording = load_recording(input_file)

        return recording

    except Exception as e:
        raise click.ClickException(
            f"Could not load '{input_file}': {e}\n"
            f"Supported formats: .sigmf, .npy, .wav, .blue\n"
            f"Use --legacy for old NPY format files"
        )


def select_params(device, sample_rate, gain, bandwidth, quiet, verbose):
    # Auto-select device if not specified
    if device is None:
        device = auto_select_tx_device(quiet)

    # Apply device-specific defaults (matching signal-testbed but conservative for TX)
    if sample_rate is None:
        # TX sample rate defaults (same as RX)
        device_sample_rates = {
            "pluto": 20e6,  # PlutoSDR up to 61 MHz, 20 MHz safe
            "hackrf": 20e6,  # HackRF up to 20 MHz
            "bladerf": 40e6,  # BladeRF up to 61 MHz, 40 MHz safe
            "usrp": 50e6,  # USRP up to 200 MHz, 50 MHz default
        }
        sample_rate = device_sample_rates.get(device, 20e6)

    if gain is None:
        # TX gain defaults (conservative for ISM band to avoid interference)
        default_tx_gains = {
            "pluto": -20,  # PlutoSDR: -20 dB (safe, low power)
            "hackrf": 0,  # HackRF: 0 dB (moderate)
            "bladerf": -10,  # BladeRF: -10 dB (conservative)
            "usrp": -10,  # USRP: -10 dB (conservative)
        }
        gain = default_tx_gains.get(device, -10)
        echo_verbose(f"Using default TX gain: {gain} dB for {device}", verbose)

    if bandwidth is None:
        # Bandwidth defaults (match sample rate)
        device_bandwidths = {
            "pluto": sample_rate,
            "hackrf": sample_rate,
            "bladerf": sample_rate,
            "usrp": sample_rate,
        }
        bandwidth = device_bandwidths.get(device)

    return device, sample_rate, gain, bandwidth


def validate_tx_gain(device_type: str, gain: float) -> None:
    """
    Validate TX gain is within device limits and warn if at extremes.

    Args:
        device_type: Type of device
        gain: TX gain in dB

    Raises:
        click.ClickException: If gain is out of range
    """
    gain_ranges = {
        "pluto": (-89, 0),
        "hackrf": (0, 47),
        "bladerf": (-15, 60),
        "usrp": (-30, 20),  # Approximate, varies by model
    }

    if device_type in gain_ranges:
        min_gain, max_gain = gain_ranges[device_type]

        if gain < min_gain or gain > max_gain:
            raise click.ClickException(
                f"TX gain {gain} dB is out of range for {device_type}\n" f"Valid range: {min_gain} to {max_gain} dB"
            )

        # Warn if at maximum
        if gain >= max_gain - 3:
            click.echo(
                click.style("WARNING: ", fg="yellow", bold=True) + f"Transmitting at high gain level ({gain} dB)\n"
                f"Maximum for {device_type}: {max_gain} dB",
                err=True,
            )


def generate_recording(generate, input_file, sample_rate, verbose, legacy):
    # Generate signal or load from file
    if generate or input_file is None:
        # Generate signal instead of loading from file
        from ria_toolkit_oss.signal.basic_signal_generator import (
            chirp,
            lfm_chirp_complex,
            sine,
            square,
        )

        # Calculate number of samples for signal generation (default: 0.1 second = 100ms)
        # Shorter duration to avoid buffer issues with large sample rates
        num_samples = int(sample_rate * 0.1)  # 100ms of signal

        if generate == "lfm" or (generate is None and input_file is None):
            # Generate LFM chirp (default - visible on spectrogram)
            echo_verbose("Generating LFM chirp signal...", verbose)
            recording = lfm_chirp_complex(
                sample_rate=int(sample_rate),
                width=int(sample_rate * 0.4),  # 40% of sample rate (safe for filter)
                chirp_period=0.001,  # 1ms chirp period
                sigfc=0,  # Baseband
                total_time=num_samples / sample_rate,
                chirp_type="up",
            )
            echo_verbose(f"Generated {len(recording.data)} sample LFM chirp", verbose)

        elif generate == "chirp":
            # Generate simple chirp
            echo_verbose("Generating chirp signal...", verbose)
            recording = chirp(sample_rate=int(sample_rate), num_samples=num_samples, center_frequency=0)  # Baseband
            echo_verbose(f"Generated {len(recording.data)} sample chirp", verbose)

        elif generate == "sine":
            # Generate sine wave at 10% offset from center
            echo_verbose("Generating sine wave signal...", verbose)
            recording = sine(
                sample_rate=int(sample_rate),
                length=num_samples,
                frequency=sample_rate * 0.1,  # 10% offset
                amplitude=0.8,
            )
            echo_verbose(f"Generated {len(recording.data)} sample sine wave", verbose)

        elif generate == "pulse":
            # Generate pulse using square wave
            echo_verbose("Generating pulse signal...", verbose)
            recording = square(
                sample_rate=int(sample_rate),
                length=num_samples,
                frequency=1000,  # 1 kHz pulse
                amplitude=0.8,
                duty_cycle=0.1,  # 10% duty cycle for pulse
            )
            echo_verbose(f"Generated {len(recording.data)} sample pulse", verbose)

        return recording

    elif input_file:
        # Load input file
        return load_input_file(input_file, legacy=legacy)

    else:
        raise click.ClickException("Either --input or --generate must be specified")


def check_sample_rate_mismatch(recording: Recording, specified_rate: float, quiet: bool) -> None:
    """
    Check if recording sample rate differs from specified rate.

    Args:
        recording: Recording object
        specified_rate: Specified sample rate
        quiet: Suppress warnings
    """
    if hasattr(recording, "metadata") and recording.metadata:
        recorded_rate = recording.metadata.get("sample_rate")
        if recorded_rate and abs(recorded_rate - specified_rate) > 1:
            if not quiet:
                click.echo(
                    click.style("Warning: ", fg="yellow")
                    + f"Recording sample rate ({format_sample_rate(recorded_rate)}) differs "
                    f"from specified rate ({format_sample_rate(specified_rate)})\n"
                    f"Using specified rate. Signal may be distorted.",
                    err=True,
                )


def repeated_transmission(sdr, recording, repeat, tx_delay, quiet, verbose):
    for i in range(repeat):
        if repeat > 1:
            echo_progress(f"\nTransmission {i + 1}/{repeat}...", quiet)

        sdr.tx_recording(recording)

        if repeat > 1:
            echo_progress(f"Transmission {i + 1}/{repeat} complete.", quiet)

        # Delay between transmissions
        if i < repeat - 1 and tx_delay > 0:
            echo_verbose(f"Waiting {tx_delay}s before next transmission...", verbose)
            time.sleep(tx_delay)

    if repeat > 1:
        echo_progress(f"\nAll {repeat} transmissions complete.", quiet)


@click.command()
@click.option("--device", "-d", type=click.Choice(TX_CAPABLE_DEVICES), help="Device type (TX-capable only)")
@click.option("--ident", "-i", help="Device identifier (IP address or name=value, e.g., 192.168.2.1 or name=myb210)")
@click.option(
    "--config", "-c", "config_file", type=click.Path(exists=True), help="Load parameters from YAML config file"
)
@click.option(
    "--sample-rate", "-s", type=float, default=None, help="Sample rate in Hz (e.g., 2e6) [default: device-specific]"
)
@click.option(
    "--center-frequency",
    "-f",
    type=str,
    default="2440M",
    show_default=True,
    help="Center frequency (e.g., 915e6, 2.4G)",
)
@click.option("--gain", "-g", type=float, help="TX gain in dB [default: device-specific safe level]")
@click.option("--bandwidth", "-b", type=float, help="Bandwidth in Hz (if supported) [default: device-specific]")
@click.option(
    "--input",
    "-in",
    "input_file",
    type=click.Path(),
    help=(
        "Input recording file (auto-detects format). "
        "If omitted and --generate not specified, generates default LFM chirp."
    ),
)
@click.option("--legacy", is_flag=True, help="Use legacy NPY format loader")
@click.option(
    "--generate",
    type=click.Choice(["lfm", "chirp", "sine", "pulse"]),
    help="Generate signal instead of loading from file (overrides --input)",
)
@click.option("--repeat", "-r", type=int, default=1, help="Repeat transmission N times (default: 1)")
@click.option("--continuous", is_flag=True, help="Transmit continuously until Ctrl+C")
@click.option("--tx-delay", type=float, default=0, help="Delay between transmissions in seconds")
@click.option("--yes", "-y", is_flag=True, help="Skip safety confirmations")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress progress output")
def transmit(
    device,
    ident,
    config_file,
    sample_rate,
    center_frequency,
    gain,
    bandwidth,
    input_file,
    legacy,
    generate,
    repeat,
    continuous,
    tx_delay,
    yes,
    verbose,
    quiet,
):
    """Transmit IQ samples from file using SDR device.

    \b
    Examples:
        ria transmit -d hackrf --generate lfm --continuous
        ria transmit -d pluto -f 2.44G -g -10 -in recordings/rec_HackRF_2MHz_2025-12-01_15-36-21_80fc33f.sigmf-data

    """

    # Load config file if specified
    config = {}
    if config_file:
        config = load_yaml_config(config_file)
        echo_verbose(f"Loaded config from: {config_file}", verbose)

    # Command-line args override config file
    device = device or config.get("device")
    ident = ident or config.get("ident") or config.get("serial")  # Support legacy 'serial' in config
    sample_rate = sample_rate or config.get("sample_rate")
    center_frequency = center_frequency or config.get("center_frequency")
    gain = gain if gain is not None else config.get("gain")
    bandwidth = bandwidth or config.get("bandwidth")
    input_file = input_file or config.get("input")
    generate = generate or config.get("generate")
    repeat = repeat if repeat != 1 else config.get("repeat", 1)
    continuous = continuous or config.get("continuous", False)
    tx_delay = tx_delay or config.get("tx_delay", 0)

    device, sample_rate, gain, bandwidth = select_params(device, sample_rate, gain, bandwidth, quiet, verbose)

    # Parse frequency
    center_freq_hz = parse_frequency(center_frequency)

    # Validate TX gain
    validate_tx_gain(device, gain)

    # Generate signal or load from file
    recording = generate_recording(generate, input_file, sample_rate, verbose, legacy)
    # Check sample rate mismatch
    check_sample_rate_mismatch(recording, sample_rate, quiet)

    # Safety warnings for continuous mode
    if continuous and not yes:
        click.echo(
            click.style("WARNING: ", fg="red", bold=True) + "Continuous transmission mode enabled\n"
            "This will transmit indefinitely until stopped.\n"
            "Ensure proper cooling and monitoring.",
            err=True,
        )
        if not click.confirm("Continue?", default=False):
            click.echo("Transmission cancelled.")
            return

    # Show transmission parameters
    num_samples = len(recording.data[0]) if len(recording.data.shape) > 1 else len(recording.data)
    echo_progress(f"Transmitting from {device.upper()}...", quiet)
    echo_progress(f"Sample rate: {format_sample_rate(sample_rate)}", quiet)
    echo_progress(f"Center frequency: {format_frequency(center_freq_hz)}", quiet)
    echo_progress(f"TX gain: {gain} dB", quiet)
    if bandwidth:
        echo_progress(f"Bandwidth: {format_sample_rate(bandwidth)}", quiet)

    # Show signal source
    if input_file:
        echo_progress(f"Input: {os.path.basename(input_file)} ({num_samples} samples)", quiet)
    else:
        signal_type = generate if generate else "lfm"
        echo_progress(f"Signal: Generated {signal_type.upper()} ({num_samples} samples)", quiet)

    if continuous:
        echo_progress("Mode: Continuous (Ctrl+C to stop)", quiet)
    elif repeat > 1:
        echo_progress(f"Repeat: {repeat} times with {tx_delay}s delay", quiet)

    # Initialize device
    echo_verbose("Initializing TX device...", verbose)
    sdr = get_sdr_device(device, ident, True)

    # Set up Ctrl+C handler for continuous mode
    stop_transmission = False

    def signal_handler(sig, frame):
        nonlocal stop_transmission
        stop_transmission = True
        click.echo("\n\nStopping transmission...")

    if continuous:
        signal.signal(signal.SIGINT, signal_handler)

    try:
        # Initialize TX with parameters
        sdr.init_tx(
            sample_rate=sample_rate, center_frequency=center_freq_hz, gain=gain, channel=0  # Default to channel 0
        )

        # Set bandwidth if supported (after init_tx)
        if bandwidth is not None and hasattr(sdr, "set_tx_bandwidth"):
            sdr.set_tx_bandwidth(bandwidth)

        # Transmission loop
        if continuous:
            echo_progress("\nTransmitting continuously... [Press Ctrl+C to stop]", quiet)

            transmission_count = 0
            while not stop_transmission:
                sdr.tx_recording(recording)
                transmission_count += 1

                if verbose and transmission_count % 10 == 0:
                    echo_verbose(f"Transmitted {transmission_count} times", verbose)

            echo_progress(f"\nTransmitted {transmission_count} times total", quiet)

        else:
            # Repeat mode or single transmission
            repeated_transmission(sdr, recording, repeat, tx_delay, quiet, verbose)

    finally:
        # Clean up device
        echo_verbose("Closing TX device...", verbose)
        if hasattr(sdr, "close"):
            sdr.close()

    echo_progress("Transmission complete!", quiet)