"""Transform command - Apply signal transformations to recordings."""

import importlib
import importlib.util
import inspect
import os
from pathlib import Path

import click

from utils.data.recording import Recording
from utils.io.recording import load_recording
from utils.transforms import iq_augmentations, iq_channel_models, iq_impairments
from utils_cli.utils.common import (
    echo_progress,
    echo_verbose,
    format_sample_count,
    save_recording,
)


def get_available_transforms(module):
    """Get list of public transform functions from a module.

    Args:
        module: Python module to inspect

    Returns:
        dict: {name: function} for all public callables
    """
    transforms = {}
    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if not name.startswith("_"):
            transforms[name] = obj
    return transforms


def get_transform_help(func):
    """Extract help info from a transform function.

    Args:
        func: Transform function to inspect

    Returns:
        dict: {description, params}
    """
    sig = inspect.signature(func)
    doc = inspect.getdoc(func) or ""

    # Get first line of docstring as description
    description = doc.split("\n")[0] if doc else "No description"

    # Extract parameters from signature (skip 'signal')
    params = {}
    for param_name, param in sig.parameters.items():
        if param_name == "signal":
            continue

        default = param.default
        param_type = "optional" if default != inspect.Parameter.empty else "required"
        default_str = f" (default: {default})" if default != inspect.Parameter.empty else ""

        params[param_name] = {
            "type": param_type,
            "default": default,
            "annotation": str(param.annotation) if param.annotation != inspect.Parameter.empty else "any",
            "display": f"{param_name} ({param_type}){default_str}",
        }

    return {"description": description, "full_doc": doc, "params": params}


def show_transform_help(transform_name, func):
    """Display compact help for a specific transform."""
    info = get_transform_help(func)

    click.echo(f"\n{transform_name}")
    click.echo("-" * 50)
    click.echo(info["description"])

    if info["params"]:
        click.echo("\nParameters:")
        for param_name, param_info in sorted(info["params"].items()):
            click.echo(f"  {param_name:20} {param_info['display']}")

    click.echo()


def quick_view_transform(recording, output_path, title="Transform Result"):
    """Create a quick PNG visualization of transformed recording using constellation plot."""
    try:
        from utils.view.view_signal_simple import view_simple_sig

        # Create PNG in same directory as output
        output_dir = Path(output_path).parent
        base_name = Path(output_path).stem
        png_path = output_dir / f"{base_name}_preview.png"

        # Use simple view with constellation
        view_simple_sig(recording, output_path=str(png_path), constellation_mode=True, title=title, saveplot=True)

        click.echo(f"Visualization saved to: {png_path}")
    except Exception as e:
        click.echo(f"Warning: Could not create visualization: {e}")


def generate_transform_suffix(transform_name, params):
    """Generate a short suffix for the output filename based on transform and params.

    Args:
        transform_name: Name of the transform
        params: Dict of parameters

    Returns:
        str: A short suffix like "awgn15" or "freqoffset10k"
    """
    suffix = transform_name.replace("_", "")

    # Add key parameter values
    if "snr_db" in params:
        suffix += f"{int(params['snr_db'])}"
    elif "snr" in params:
        suffix += f"{int(params['snr'])}"
    elif "amplitude_variance" in params:
        suffix += f"{int(params['amplitude_variance']*100)}av"
    elif "phase_variance" in params:
        suffix += f"{int(params['phase_variance']*100000)}pv"
    elif "compression_gain" in params:
        suffix += f"{params['compression_gain']:.2f}".rstrip("0").rstrip(".")
    elif "offset_hz" in params:
        hz = params["offset_hz"]
        if abs(hz) >= 1e6:
            suffix += f"{hz/1e6:.0f}m"
        elif abs(hz) >= 1e3:
            suffix += f"{hz/1e3:.0f}k"
        else:
            suffix += f"{hz:.0f}"
    elif "offset" in params:
        suffix += f"{params['offset']:.2f}".rstrip("0").rstrip(".")
    elif "doppler_hz" in params:
        suffix += f"{params['doppler_hz']:.0f}"

    return suffix


def parse_transform_params(param_strings):
    """Parse transform parameters from CLI options.

    Args:
        param_strings: List of 'KEY=VALUE' strings

    Returns:
        dict: {key: value} with types inferred
    """
    params = {}
    if not param_strings:
        return params

    for param_str in param_strings:
        if "=" not in param_str:
            raise click.BadParameter(f"Parameter must be KEY=VALUE, got: {param_str}")

        key, value = param_str.split("=", 1)
        key = key.strip()
        value = value.strip()

        # Try to infer type
        try:
            # Try to parse scientific notation and floats
            if "e" in value.lower() or "." in value:
                params[key] = float(value)
            else:
                params[key] = int(value)
        except ValueError:
            # Keep as string
            params[key] = value

    return params


def load_custom_transforms(transform_dir):
    """Load custom transform functions from a directory.

    Args:
        transform_dir: Path to directory containing .py files with transform functions

    Returns:
        dict: {transform_name: function} for all public functions in all .py files

    Raises:
        click.ClickException: If directory doesn't exist or no transforms found
    """
    transform_dir = Path(transform_dir)

    if not transform_dir.exists():
        raise click.ClickException(f"Transform directory does not exist: {transform_dir}")

    if not transform_dir.is_dir():
        raise click.ClickException(f"Path is not a directory: {transform_dir}")

    transforms = {}
    py_files = list(transform_dir.glob("*.py"))

    if not py_files:
        raise click.ClickException(f"No .py files found in {transform_dir}")

    for py_file in py_files:
        try:
            # Load module dynamically
            spec = importlib.util.spec_from_file_location(py_file.stem, py_file)
            if spec is None or spec.loader is None:
                click.echo(f"Warning: Could not load {py_file.name}")
                continue

            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)

            # Extract all public functions
            for name, obj in inspect.getmembers(module, inspect.isfunction):
                if not name.startswith("_"):
                    # Store with source file info for metadata
                    obj._transform_source_file = py_file.name
                    transforms[name] = obj
        except Exception as e:
            raise click.ClickException(f"Failed to load {py_file.name}: {e}")

    return transforms


def check_input_errors(item_name: str, item, available, input, help_transform):
    if item is None:
        if help_transform:
            raise click.UsageError(f"{item_name.upper()} must be specified for --help-transform")
        else:
            raise click.UsageError(f"{item_name.upper()} must be specified (or use --list)")
    if item not in available:
        raise click.ClickException(f"Unknown {item_name}: {item}\n" f"Use --list to see available options")
    if input is None and not help_transform:
        raise click.UsageError("INPUT must be specified")


def load_input(input, verbose):
    # Load input
    try:
        recording = load_recording(input)
    except Exception as e:
        raise click.ClickException(f"Failed to load input: {e}")

    echo_verbose(f"Loaded {format_sample_count(recording.data.shape[-1])} samples", verbose)
    return recording

@click.group()
def transform():
    """Apply signal transformations to recordings.

    Transform supports three categories of operations:
    - augment: Modify signal to create new ML examples
    - impair: Degrade signal with noise, distortion, etc.
    - apply_channel: Apply channel models (fading, Doppler, etc.)

    Each operation is applied independently. Chain multiple transforms by
    running this command multiple times.
    
    \b
    Examples:
        # List available augmentations 
        utils transform augment --list
    \b
        # Apply channel swap
        utils transform augment channel_swap input.npy
    \b
        # Apply AWGN impairment
        utils transform impair awgn input.npy --snr-db 15
    \b
        # Apply Rayleigh fading channel
        utils transform apply_channel rayleigh input.npy --num-paths 5
    """
    pass


@transform.command(name="augment")
@click.argument("augmentation", required=False)
@click.argument("input", type=click.Path(exists=True), required=False)
@click.argument("output", type=click.Path(), required=False)
@click.option("--list", "list_transforms", is_flag=True, help="List available augmentations")
@click.option("--help-transform", is_flag=True, help="Show parameters for this augmentation")
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
def augment(augmentation, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet):
    """Apply augmentation transforms to recordings.

    Augmentations modify signals to create new training examples without
    degrading quality (e.g., channel swap, time reversal, quantization).

    Examples:

      # List all augmentations
      \b
      utils transform augment --list

      # Show parameters for an augmentation
      \b
      utils transform augment channel_swap --help-transform

      # Apply augmentation
      \b
      utils transform augment channel_swap input.npy

      # Apply with parameters and save visualization
      \b
      utils transform augment drop_samples input.npy --params max_section_size=5 --view
    """
    available = get_available_transforms(iq_augmentations)

    if list_transforms:
        click.echo("Available augmentations:")
        for name in sorted(available.keys()):
            func = available[name]
            docstring = (func.__doc__ or "").split("\n")[0].strip()
            click.echo(f"  {name:30} {docstring}")
        return

    if help_transform:
        check_input_errors("augmentation", augmentation, available, input, help_transform)
        show_transform_help(augmentation, available[augmentation])
        return

    check_input_errors("augmentation", augmentation, available, input, help_transform)

    # Generate output filename if not provided
    if output is None:
        input_path = Path(input)
        input_stem = input_path.stem
        ext = input_path.suffix
        suffix = generate_transform_suffix(augmentation, parse_transform_params(params))
        output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
        echo_verbose(f"Auto-generated output: {output}", verbose)

    # Check if output exists
    if not overwrite and Path(output).exists():
        raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")

    echo_progress(f"Augmenting: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
    echo_verbose(f"Transform: {augmentation}", verbose)

    # Load input
    recording = load_input(input, verbose)

    # Parse and apply transform
    try:
        transform_func = available[augmentation]
        transform_params = parse_transform_params(params)
        echo_verbose(f"Parameters: {transform_params}", verbose)

        result = transform_func(recording, **transform_params)
    except Exception as e:
        raise click.ClickException(f"Transform failed: {e}")

    # Track transform in metadata (Recording.metadata is a property that returns a copy)
    # So we need to work with a copy and create a new Recording with updated metadata
    updated_metadata = result.metadata.copy()
    if "transforms_applied" not in updated_metadata:
        updated_metadata["transforms_applied"] = []

    updated_metadata["transforms_applied"].append(
        {"type": "augment", "name": augmentation, "params": parse_transform_params(params)}
    )

    # Create new recording with updated metadata
    result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)

    # Save output
    try:
        save_recording(result, output, overwrite=overwrite, verbose=verbose)
        echo_progress(f"Saved to: {output}", quiet)
    except Exception as e:
        raise click.ClickException(f"Failed to save output: {e}")

    # Optional: Create visualization
    if view:
        echo_verbose("Creating visualization...", verbose)
        quick_view_transform(result, output, title=f"{augmentation.replace('_', ' ').title()} - {Path(output).name}")


@transform.command(name="impair")
@click.argument("impairment", required=False)
@click.argument("input", type=click.Path(exists=True), required=False)
@click.argument("output", type=click.Path(), required=False)
@click.option("--list", "list_transforms", is_flag=True, help="List available impairments")
@click.option("--help-transform", is_flag=True, help="Show parameters for this impairment")
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
def impair(impairment, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet):
    """Apply impairment transforms to recordings.

    Impairments degrade signals by adding noise, distortion, and other
    channel effects (e.g., AWGN, phase noise, IQ imbalance).

    Examples:

      # List all impairments
      \b
      utils transform impair --list

      # Show parameters for an impairment
      \b
      utils transform impair add_awgn_to_signal --help-transform

      # Apply impairment
      \b
      utils transform impair add_awgn_to_signal input.npy --params snr=10

      # Apply with visualization
      \b
      utils transform impair add_phase_noise input.npy --params phase_variance=0.001 --view
    """
    available = get_available_transforms(iq_impairments)

    if list_transforms:
        click.echo("Available impairments:")
        for name in sorted(available.keys()):
            func = available[name]
            docstring = (func.__doc__ or "").split("\n")[0].strip()
            click.echo(f"  {name:30} {docstring}")
        return

    if help_transform:
        check_input_errors("impairment", impairment, available, input, help_transform)
        show_transform_help(impairment, available[impairment])
        return

    check_input_errors("impairment", impairment, available, input, help_transform)

    # Generate output filename if not provided
    if output is None:
        input_path = Path(input)
        input_stem = input_path.stem
        ext = input_path.suffix
        suffix = generate_transform_suffix(impairment, parse_transform_params(params))
        output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
        echo_verbose(f"Auto-generated output: {output}", verbose)

    # Check if output exists
    if not overwrite and Path(output).exists():
        raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")

    echo_progress(f"Impairing: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
    echo_verbose(f"Transform: {impairment}", verbose)

    # Load input
    recording = load_input(input, verbose)

    # Parse and apply transform
    try:
        transform_func = available[impairment]
        transform_params = parse_transform_params(params)
        echo_verbose(f"Parameters: {transform_params}", verbose)

        result = transform_func(recording, **transform_params)
    except Exception as e:
        raise click.ClickException(f"Transform failed: {e}")

    # Track transform in metadata (Recording.metadata is a property that returns a copy)
    updated_metadata = result.metadata.copy()
    if "transforms_applied" not in updated_metadata:
        updated_metadata["transforms_applied"] = []

    updated_metadata["transforms_applied"].append(
        {"type": "impair", "name": impairment, "params": parse_transform_params(params)}
    )

    # Create new recording with updated metadata
    result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)

    # Save output
    try:
        save_recording(result, output, overwrite=overwrite, verbose=verbose)
        echo_progress(f"Saved to: {output}", quiet)
    except Exception as e:
        raise click.ClickException(f"Failed to save output: {e}")

    # Optional: Create visualization
    if view:
        echo_verbose("Creating visualization...", verbose)
        quick_view_transform(result, output, title=f"{impairment.replace('_', ' ').title()} - {Path(output).name}")


@transform.command(name="apply_channel")
@click.argument("channel_model", required=False)
@click.argument("input", type=click.Path(exists=True), required=False)
@click.argument("output", type=click.Path(), required=False)
@click.option("--list", "list_transforms", is_flag=True, help="List available channel models")
@click.option("--help-transform", is_flag=True, help="Show parameters for this channel model")
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
def apply_channel(
    channel_model, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet
):
    """Apply channel models to recordings.

    Channel models simulate RF propagation effects like fading, Doppler shift,
    and multipath reflections.

    Use --list to see available channel models and their parameters.

    \b
    Examples:
        utils transform apply_channel rayleigh_fading_channel input.npy --params num_paths=3 snr_db=15

    \b
        utils transform apply_channel doppler_channel recordings/input.npy \\
        --params satellite_velocity=7500 \\
        --params satellite_initial_distance=400000 \\
        --params frequency=1e9 \\
        --params sample_rate=2e6
    """
    available = get_available_transforms(iq_channel_models)

    if list_transforms:
        click.echo("Available channel models:")
        for name in sorted(available.keys()):
            func = available[name]
            docstring = (func.__doc__ or "").split("\n")[0].strip()
            click.echo(f"  {name:30} {docstring}")
        return

    if help_transform:
        check_input_errors("channel_model", channel_model, available, input, help_transform)
        show_transform_help(channel_model, available[channel_model])
        return

    check_input_errors("channel_model", channel_model, available, input, help_transform)

    # Generate output filename if not provided
    if output is None:
        input_path = Path(input)
        input_stem = input_path.stem
        ext = input_path.suffix
        suffix = generate_transform_suffix(channel_model, parse_transform_params(params))
        output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
        echo_verbose(f"Auto-generated output: {output}", verbose)

    # Check if output exists
    if not overwrite and Path(output).exists():
        raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")

    echo_progress(f"Applying channel: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
    echo_verbose(f"Channel model: {channel_model}", verbose)

    # Load input
    recording = load_input(input, verbose)

    # Parse and apply transform
    try:
        transform_func = available[channel_model]
        transform_params = parse_transform_params(params)
        echo_verbose(f"Parameters: {transform_params}", verbose)

        result = transform_func(recording, **transform_params)
    except Exception as e:
        raise click.ClickException(f"Transform failed: {e}")

    # Track transform in metadata (Recording.metadata is a property that returns a copy)
    updated_metadata = result.metadata.copy()
    if "transforms_applied" not in updated_metadata:
        updated_metadata["transforms_applied"] = []

    updated_metadata["transforms_applied"].append(
        {"type": "channel", "name": channel_model, "params": parse_transform_params(params)}
    )

    # Create new recording with updated metadata
    result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)

    # Save output
    try:
        save_recording(result, output, overwrite=overwrite, verbose=verbose)
        echo_progress(f"Saved to: {output}", quiet)
    except Exception as e:
        raise click.ClickException(f"Failed to save output: {e}")

    # Optional: Create visualization
    if view:
        echo_verbose("Creating visualization...", verbose)
        quick_view_transform(result, output, title=f"{channel_model.replace('_', ' ').title()} - {Path(output).name}")


@transform.command(name="custom")
@click.argument("transform_name", required=False)
@click.argument("input", type=click.Path(exists=True), required=False)
@click.argument("output", type=click.Path(), required=False)
@click.option(
    "--transform-dir",
    type=click.Path(exists=True),
    required=True,
    help="Path to directory containing custom transform .py files",
)
@click.option("--list", "list_transforms", is_flag=True, help="List available custom transforms")
@click.option("--help-transform", is_flag=True, help="Show parameters for this transform")
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
def custom(
    transform_name,
    input,
    output,
    transform_dir,
    list_transforms,
    help_transform,
    params,
    view,
    overwrite,
    verbose,
    quiet,
):
    """Apply custom user-defined transforms to recordings.

    Custom transforms are Python functions loaded from user-specified directory.
    Each .py file in the directory is scanned for public functions that can be used.

    Transform functions must have signature:
        def my_transform(signal, **kwargs) -> signal_or_recording
    where signal is a complex CxN array or Recording object.

    Examples:

      # List all custom transforms in directory
      \b
      utils transform custom --transform-dir ~/my_transforms --list

      # Show parameters for a transform
      \b
      utils transform custom my_filter --transform-dir ~/my_transforms --help-transform

      # Apply custom transform
      \b
      utils transform custom my_filter input.npy --transform-dir ~/my_transforms

      # With parameters and visualization
      \b
      utils transform custom my_filter input.npy --transform-dir ~/my_transforms \\
        --params cutoff_freq=5000 order=4 --view
    """
    try:
        available = load_custom_transforms(transform_dir)
    except click.ClickException:
        raise

    if list_transforms:
        click.echo(f"Available custom transforms in {transform_dir}:")
        for name in sorted(available.keys()):
            func = available[name]
            source_file = getattr(func, "_transform_source_file", "unknown")
            docstring = (func.__doc__ or "").split("\n")[0].strip()
            click.echo(f"  {name:30} {docstring:40} [{source_file}]")
        return

    if help_transform:
        check_input_errors("transform_name", transform_name, available, input, help_transform)
        show_transform_help(transform_name, available[transform_name])
        return

    check_input_errors("transform_name", transform_name, available, input, help_transform)

    # Generate output filename if not provided
    if output is None:
        input_path = Path(input)
        input_stem = input_path.stem
        ext = input_path.suffix
        suffix = generate_transform_suffix(transform_name, parse_transform_params(params))
        output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
        echo_verbose(f"Auto-generated output: {output}", verbose)

    # Check if output exists
    if not overwrite and Path(output).exists():
        raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")

    echo_progress(f"Applying custom: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
    echo_verbose(f"Transform: {transform_name}", verbose)

    # Load input
    recording = load_input(input, verbose)

    # Parse and apply transform
    try:
        transform_func = available[transform_name]
        transform_params = parse_transform_params(params)
        echo_verbose(f"Parameters: {transform_params}", verbose)

        result = transform_func(recording, **transform_params)
    except Exception as e:
        raise click.ClickException(f"Transform failed: {e}")

    # Track transform in metadata
    updated_metadata = result.metadata.copy()
    if "transforms_applied" not in updated_metadata:
        updated_metadata["transforms_applied"] = []

    updated_metadata["transforms_applied"].append(
        {
            "type": "custom",
            "name": transform_name,
            "source_file": getattr(available[transform_name], "_transform_source_file", "unknown"),
            "params": parse_transform_params(params),
        }
    )

    # Create new recording with updated metadata
    result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)

    # Save output
    try:
        save_recording(result, output, overwrite=overwrite, verbose=verbose)
        echo_progress(f"Saved to: {output}", quiet)
    except Exception as e:
        raise click.ClickException(f"Failed to save output: {e}")

    # Optional: Create visualization
    if view:
        echo_verbose("Creating visualization...", verbose)
        quick_view_transform(result, output, title=f"{transform_name.replace('_', ' ').title()} - {Path(output).name}")