"""Capture command for SDR devices."""

import os

import click

from utils.io import to_blue, to_npy, to_sigmf, to_wav
from utils.io.recording import generate_filename
from utils.view.view_signal_simple import view_simple_sig

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


def list_all_devices():
    # Load drivers and collect all devices
    load_sdr_drivers(verbose=False)

    all_devices = []
    all_devices.extend(find_uhd_devices())
    all_devices.extend(find_pluto_devices())
    all_devices.extend(find_hackrf_devices())
    all_devices.extend(find_bladerf_devices())
    all_devices.extend(find_rtlsdr_devices())
    all_devices.extend(find_thinkrf_devices())

    return all_devices


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

    Args:
        quiet: Suppress warning messages

    Returns:
        Device type string

    Raises:
        click.ClickException: If no devices or multiple devices found
    """
    all_devices = list_all_devices()

    if len(all_devices) == 0:
        raise click.ClickException("No SDR devices found.\n" "Run 'utils discover' to see available devices.")

    elif len(all_devices) == 1:
        device = all_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",
            "rtlsdr": "rtlsdr",
            "thinkrf": "thinkrf",
        }

        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 all_devices)
        raise click.ClickException(
            f"Multiple devices found. Specify with --device\n\n"
            f"Available devices:\n{device_list}\n\n"
            f"Run 'utils discover' for more details."
        )


def get_metadata_dict(config, metadata):
    # Parse metadata - start with user config defaults
    metadata_dict = config.get("metadata", {})

    # Load user config and apply defaults
    user_config = load_user_config()

    # Apply user config metadata (if user config exists)
    if user_config:
        # Add standard metadata fields from config
        for key in ["author", "organization", "project", "location", "testbed"]:
            if key in user_config and key not in metadata_dict:
                metadata_dict[key] = user_config[key]

        # Add SigMF fields from config
        if "sigmf" in user_config:
            sigmf = user_config["sigmf"]
            for key in ["license", "hw", "dataset"]:
                if key in sigmf and key not in metadata_dict:
                    metadata_dict[key] = sigmf[key]

    # CLI metadata overrides everything
    if metadata:
        metadata_dict.update(parse_metadata_args(metadata))

    return metadata_dict


def save_visualization(recording, output_file: str, quiet: bool = False):
    """Save visualization of recording.

    Args:
        recording: Recording object
        output_file: Path to save visualization (PNG)
        quiet: Suppress progress messages
    """
    # Generate image filename matching recording filename
    base_name = os.path.splitext(output_file)[0]
    if output_file.endswith(".sigmf-data"):
        base_name = output_file.replace(".sigmf-data", "")
    output_file = base_name + ".png"

    try:
        echo_progress(f"Generating visualization: {output_file}", quiet)
        view_simple_sig(recording, output_path=output_file, saveplot=True, fast_mode=False, labels_mode=True)
    except ImportError as e:
        click.echo(click.style("Warning: ", fg="yellow") + f"Could not save visualization: {e}", err=True)
    except Exception as e:
        click.echo(click.style("Warning: ", fg="yellow") + f"Failed to save visualization: {e}", err=True)


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

    # Apply device-specific defaults (matching signal-testbed)
    if sample_rate is None:
        # Sample rate defaults based on signal-testbed hardware limits
        device_sample_rates = {
            "rtlsdr": 2.4e6,  # RTL-SDR max is 3.2 MHz, use 2.4 MHz safe default
            "thinkrf": 31.25e6,  # ThinkRF decimation 4 (from 125 MS/s)
            "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 from signal-testbed
        }
        sample_rate = device_sample_rates.get(device, 20e6)

    if gain is None:
        # RX gain defaults (matching signal-testbed's 32 dB baseline, adjusted per device)
        default_gains = {
            "pluto": 32,
            "hackrf": 32,
            "bladerf": 32,
            "usrp": 32,
            "rtlsdr": 32,  # RTL-SDR will auto-select closest valid gain
            "thinkrf": 0,  # ThinkRF uses attenuation, 0 = no attenuation
        }
        gain = default_gains.get(device, 32)
        echo_verbose(f"Using default RX gain: {gain} dB for {device}", verbose)

    if bandwidth is None:
        # Bandwidth defaults (match sample rate for most devices)
        device_bandwidths = {
            "rtlsdr": None,  # RTL-SDR doesn't support bandwidth setting
            "thinkrf": None,  # ThinkRF manages bandwidth internally
            "pluto": sample_rate,
            "hackrf": sample_rate,
            "bladerf": sample_rate,
            "usrp": sample_rate,
        }
        bandwidth = device_bandwidths.get(device)

    return device, sample_rate, gain, bandwidth


def determine_output_format(output, output_format, output_dir):
    # Determine output format and save
    # If output specified, parse directory and filename
    if output:
        # Auto-detect format from extension if not specified
        if output_format is None:
            ext = os.path.splitext(output)[1].lower().lstrip(".")
            if ext in ["sigmf", "sigmf-data"]:
                output_format = "sigmf"
            elif ext == "npy":
                output_format = "npy"
            elif ext == "wav":
                output_format = "wav"
            elif ext == "blue":
                output_format = "blue"
            else:
                # Default to SigMF
                output_format = "sigmf"

        # Get output directory and filename from provided path
        output_path_dir = os.path.dirname(output)
        if output_path_dir:
            output_dir = output_path_dir
        output_filename = os.path.basename(output)

        # Remove extension for formats that add it
        if output_format == "sigmf":
            output_filename = output_filename.replace(".sigmf-data", "").replace(".sigmf", "")
    else:
        # Use auto-generated filename based on timestamp and rec_id
        output_filename = None  # Will be auto-generated by save functions
        if output_format is None:
            output_format = "sigmf"  # Default format

    return output_format, output_filename, output_dir


# ============================================================================
# Main command
# ============================================================================

@click.command()
@click.argument("inputs", nargs=-1, required=True, type=click.Path(exists=True))
@click.argument("output", nargs=1, required=True, type=click.Path())
@click.option(
    "--device",
    "-d",
    type=click.Choice(["pluto", "hackrf", "bladerf", "usrp", "rtlsdr", "thinkrf"]),
    help="Device type",
)
@click.option("--ident", "-i", help="Device identifier (IP address or name=value, e.g., 192.168.2.1 or name=mypluto)")
@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="RX gain in dB [default: device-specific]")
@click.option("--bandwidth", "-b", type=float, help="Bandwidth in Hz (if supported) [default: device-specific]")
@click.option("--num-samples", "-n", type=int, show_default=True, help="Number of samples to capture")
@click.option("--duration", "-t", type=float, help="Duration in seconds (alternative to --num-samples)")
@click.option("--output", "-o", help="Output filename (defaults to auto-generated with timestamp)")
@click.option("--output-dir", default="recordings", help="Output directory (default: recordings/)")
@click.option(
    "--format",
    "output_format",
    type=click.Choice(["npy", "sigmf", "wav", "blue"]),
    help="Output format (default: sigmf)",
)
@click.option("--save-image", is_flag=True, help="Save visualization PNG alongside recording")
@click.option("--metadata", "-m", multiple=True, help="Add custom metadata (KEY=VALUE)")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress progress output")
def capture(
    device,
    ident,
    config_file,
    sample_rate,
    center_frequency,
    gain,
    bandwidth,
    num_samples,
    duration,
    output,
    output_dir,
    output_format,
    save_image,
    metadata,
    verbose,
    quiet,
):
    """Capture IQ samples from SDR device and save to file.

    \b
    Examples:
        utils capture -d hackrf -s 2e6 -f 2.44e6 -b 2e6
        utils capture -d pluto -s 1e6 -f 2e9 -b 2e6 -n 50
    
    """

    # 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 or config.get("gain")
    bandwidth = bandwidth or config.get("bandwidth")
    num_samples = num_samples or config.get("num_samples")
    duration = duration or config.get("duration")
    output = output or config.get("output")
    output_format = output_format or config.get("format")

    # Parse metadata
    metadata_dict = get_metadata_dict(config=config, metadata=metadata)

    # Select parameters
    device, sample_rate, gain, bandwidth = select_params(
        device=device, sample_rate=sample_rate, gain=gain, bandwidth=bandwidth, quiet=quiet, verbose=verbose
    )

    # Parse frequency
    center_freq_hz = parse_frequency(center_frequency)

    # Calculate num_samples from duration if needed
    if duration is not None and num_samples is None:
        num_samples = int(duration * sample_rate)
        echo_verbose(f"Duration {duration}s = {num_samples} samples at {format_sample_rate(sample_rate)}", verbose)

    # Show capture parameters
    echo_progress(f"Capturing 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)
    if gain is not None:
        echo_progress(f"Gain: {gain} dB", quiet)
    if bandwidth is not None:
        echo_progress(f"Bandwidth: {format_sample_rate(bandwidth)}", quiet)

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

    try:
        # Initialize RX with parameters
        echo_verbose("Initializing RX...", verbose)
        sdr.init_rx(
            sample_rate=sample_rate, center_frequency=center_freq_hz, gain=gain, channel=0  # Default to channel 0
        )

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

        # Capture
        echo_progress(f"Capturing {num_samples} samples...", quiet)
        recording = sdr.record(num_samples=num_samples)

        echo_progress(
            f"Captured {recording.data.shape[1] if len(recording.data.shape) > 1 else len(recording.data)} samples",
            quiet,
        )

        # Add custom metadata to recording
        if metadata_dict:
            for key, value in metadata_dict.items():
                recording.update_metadata(key, value)

        output_format, output_filename, output_dir = determine_output_format(
            output=output, output_format=output_format, output_dir=output_dir
        )
        echo_progress(f"Saving to {output_format.upper()} format...", quiet)

        # Save recording (filenames with timestamp auto-generated if output_filename is None)
        # All to_* functions handle directory creation internally
        # Note: to_sigmf returns None, others return path
        if output_format == "sigmf":
            to_sigmf(recording, filename=output_filename, path=output_dir)
            # Build path manually since to_sigmf doesn't return it
            base_name = (
                os.path.splitext(output_filename)[0] if output_filename else generate_filename(recording=recording)
            )
            saved_path = os.path.join(output_dir, f"{base_name}.sigmf-data")
        elif output_format == "npy":
            saved_path = to_npy(recording, filename=output_filename, path=output_dir)
        elif output_format == "wav":
            saved_path = to_wav(recording, filename=output_filename, path=output_dir)
        elif output_format == "blue":
            saved_path = to_blue(recording, filename=output_filename, path=output_dir)

        echo_progress(f"Saved to: {saved_path}", quiet)

        # Save visualization if requested
        if save_image:
            save_visualization(recording, saved_path, quiet)

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

    echo_progress("Capture complete!", quiet)