help_probe/

completion.rs

use crate::model::ProbeResult;

/// Clean a flag name by removing trailing `...` (used for variadic flags).
/// This ensures flags are valid in all shell completion formats.
fn clean_flag_name(flag: &str) -> String {
    flag.trim_end_matches("...").to_string()
}

/// Clean a subcommand name by removing trailing commas and trimming.
/// Also filters out invalid entries like `...`.
fn clean_subcommand_name(name: &str) -> Option<String> {
    let cleaned = name.trim_end_matches(',').trim().to_string();
    if cleaned.is_empty() || cleaned == "..." {
        None
    } else {
        Some(cleaned)
    }
}

/// Collect and clean all subcommands from a ProbeResult.
/// Returns a vector of cleaned subcommand names.
fn collect_clean_subcommands(result: &ProbeResult) -> Vec<String> {
    result
        .subcommands
        .iter()
        .filter_map(|s| clean_subcommand_name(&s.name))
        .collect()
}

/// Shell types for which we can generate completions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Shell {
    /// Bash shell
    Bash,
    /// Zsh shell
    Zsh,
    /// Fish shell
    Fish,
    /// PowerShell
    PowerShell,
    /// NuShell
    NuShell,
}

impl Shell {
    /// Parse shell name from string (case-insensitive).
    ///
    /// # Examples
    ///
    /// ```
    /// use help_probe::completion::Shell;
    ///
    /// assert_eq!(Shell::from_str("bash"), Some(Shell::Bash));
    /// assert_eq!(Shell::from_str("BASH"), Some(Shell::Bash));
    /// assert_eq!(Shell::from_str("zsh"), Some(Shell::Zsh));
    /// assert_eq!(Shell::from_str("powershell"), Some(Shell::PowerShell));
    /// assert_eq!(Shell::from_str("pwsh"), Some(Shell::PowerShell));
    /// assert_eq!(Shell::from_str("nushell"), Some(Shell::NuShell));
    /// assert_eq!(Shell::from_str("nu"), Some(Shell::NuShell));
    /// assert_eq!(Shell::from_str("invalid"), None);
    /// ```
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "bash" => Some(Shell::Bash),
            "zsh" => Some(Shell::Zsh),
            "fish" => Some(Shell::Fish),
            "powershell" | "pwsh" => Some(Shell::PowerShell),
            "nushell" | "nu" => Some(Shell::NuShell),
            _ => None,
        }
    }
}

/// Generate shell completion script for a command.
///
/// This generates a complete completion script that can be sourced
/// to provide tab completion for the probed command.
///
/// # Examples
///
/// ```
/// use help_probe::{completion::{generate_shell_completion, Shell}, model::ProbeResult};
///
/// let result = ProbeResult {
///     command: "mytool".to_string(),
///     args: vec![],
///     exit_code: Some(0),
///     timed_out: false,
///     help_flag_detected: true,
///     usage_blocks: vec![],
///     options: vec![],
///     subcommands: vec![],
///     arguments: vec![],
///     examples: vec![],
///     environment_variables: vec![],
///     validation_rules: vec![],
///     raw_stdout: String::new(),
///     raw_stderr: String::new(),
/// };
///
/// let bash_completion = generate_shell_completion(&result, Shell::Bash);
/// assert!(bash_completion.contains("mytool"));
/// assert!(bash_completion.contains("complete -F"));
/// ```
pub fn generate_shell_completion(result: &ProbeResult, shell: Shell) -> String {
    match shell {
        Shell::Bash => generate_bash_completion(result),
        Shell::Zsh => generate_zsh_completion(result),
        Shell::Fish => generate_fish_completion(result),
        Shell::PowerShell => generate_powershell_completion(result),
        Shell::NuShell => generate_nushell_completion(result),
    }
}

/// Generate bash completion script.
fn generate_bash_completion(result: &ProbeResult) -> String {
    let cmd_name = &result.command;
    let func_name = format!(
        "_{}_completion",
        cmd_name.replace('-', "_").replace('/', "_")
    );

    let mut script = String::new();
    script.push_str(&format!("# Bash completion for {}\n", cmd_name));
    script.push_str(&format!("# Generated by help-probe\n\n"));

    script.push_str(&format!("{}() {{\n", func_name));
    script.push_str("    local cur prev words cword\n");
    script.push_str("    COMPREPLY=()\n");
    script.push_str("    cur=\"${COMP_WORDS[COMP_CWORD]}\"\n");
    script.push_str("    prev=\"${COMP_WORDS[COMP_CWORD-1]}\"\n");
    script.push_str("    words=(\"${COMP_WORDS[@]}\")\n");
    script.push_str("    cword=$COMP_CWORD\n\n");

    // Collect all options (clean trailing ... from flags)
    let mut all_options: Vec<String> = Vec::new();
    for opt in &result.options {
        all_options.extend(opt.short_flags.iter().map(|f| clean_flag_name(f)));
        all_options.extend(opt.long_flags.iter().map(|f| clean_flag_name(f)));
    }

    // Collect all subcommands (clean trailing commas)
    let subcommands = collect_clean_subcommands(result);

    script.push_str("    # Options\n");
    script.push_str(&format!(
        "    local opts=({})\n",
        all_options
            .iter()
            .map(|o| format!("\"{}\"", o))
            .collect::<Vec<_>>()
            .join(" ")
    ));

    if !subcommands.is_empty() {
        script.push_str("    # Subcommands\n");
        script.push_str(&format!(
            "    local subcommands=({})\n",
            subcommands
                .iter()
                .map(|s| format!("\"{}\"", s))
                .collect::<Vec<_>>()
                .join(" ")
        ));
    }

    script.push_str("\n");
    script.push_str("    # Complete options and subcommands\n");
    script.push_str("    if [[ \"$cur\" == -* ]]; then\n");
    script.push_str("        COMPREPLY=($(compgen -W \"${opts[*]}\" -- \"$cur\"))\n");
    script.push_str("    elif [[ ${#words[@]} -eq 2 ]]; then\n");
    if !subcommands.is_empty() {
        script.push_str("        COMPREPLY=($(compgen -W \"${subcommands[*]}\" -- \"$cur\"))\n");
    } else {
        script.push_str("        COMPREPLY=($(compgen -f -- \"$cur\"))\n");
    }
    script.push_str("    else\n");
    script.push_str("        # Complete files for arguments\n");
    script.push_str("        COMPREPLY=($(compgen -f -- \"$cur\"))\n");
    script.push_str("    fi\n");
    script.push_str("}\n\n");

    script.push_str(&format!("complete -F {} {}\n", func_name, cmd_name));

    script
}

/// Generate zsh completion script.
fn generate_zsh_completion(result: &ProbeResult) -> String {
    let cmd_name = &result.command;
    let func_name = format!(
        "_{}_completion",
        cmd_name.replace('-', "_").replace('/', "_")
    );

    let mut script = String::new();
    script.push_str(&format!("#compdef {}\n", cmd_name));
    script.push_str(&format!("# Zsh completion for {}\n", cmd_name));
    script.push_str(&format!("# Generated by help-probe\n\n"));

    // Wrap in a function to avoid issues when sourcing directly
    script.push_str(&format!("{}() {{\n", func_name));

    // Collect all options (clean trailing ... from flags)
    let mut all_options: Vec<String> = Vec::new();
    for opt in &result.options {
        all_options.extend(opt.short_flags.iter().map(|f| clean_flag_name(f)));
        all_options.extend(opt.long_flags.iter().map(|f| clean_flag_name(f)));
    }

    // Collect subcommands (clean trailing commas)
    let subcommands = collect_clean_subcommands(result);

    script.push_str("  local -a opts=(\n");
    for opt in &all_options {
        script.push_str(&format!("    '{}'\n", opt));
    }
    script.push_str("  )\n\n");

    if !subcommands.is_empty() {
        script.push_str("  local -a subcommands=(\n");
        for sub in &subcommands {
            script.push_str(&format!("    '{}'\n", sub));
        }
        script.push_str("  )\n\n");
    }

    script.push_str("  _arguments \\\n");

    // Add options
    for opt in &result.options {
        for long_flag in &opt.long_flags {
            // Clean trailing ... from flag name
            let mut arg_spec = clean_flag_name(long_flag);
            if opt.takes_argument {
                if let Some(arg_name) = &opt.argument_name {
                    arg_spec.push_str(&format!(":{}:", arg_name));
                } else {
                    arg_spec.push_str(":value:");
                }
            }
            // Clean short flags too
            let clean_short_flags: Vec<String> =
                opt.short_flags.iter().map(|f| clean_flag_name(f)).collect();
            script.push_str(&format!(
                "    '({}){}' \\\n",
                clean_short_flags.join(","),
                arg_spec
            ));
        }
    }

    if !subcommands.is_empty() {
        script.push_str(&format!("    '1: :->subcommands' \\\n"));
        script.push_str("    '*: :->files'\n\n");
        script.push_str("  case $state in\n");
        script.push_str("    subcommands)\n");
        script.push_str("      _describe 'subcommands' subcommands\n");
        script.push_str("      ;;\n");
        script.push_str("    files)\n");
        script.push_str("      _files\n");
        script.push_str("      ;;\n");
        script.push_str("  esac\n");
    } else {
        script.push_str("    '*: :_files'\n");
    }

    script.push_str("}\n\n");
    script.push_str(&format!("{}\n", func_name));

    script
}

/// Generate fish completion script.
/// Fish completions should have each option as a separate `complete` command.
fn generate_fish_completion(result: &ProbeResult) -> String {
    let cmd_name = &result.command;

    let mut script = String::new();
    script.push_str(&format!("# Fish completion for {}\n", cmd_name));
    script.push_str(&format!("# Generated by help-probe\n\n"));

    // Add options - each option gets its own complete command
    for opt in &result.options {
        // Process long flags
        for long_flag in &opt.long_flags {
            // Strip -- prefix and clean trailing ... (Fish doesn't allow ... in flag names)
            let mut flag_name = long_flag.trim_start_matches("--").to_string();
            flag_name = clean_flag_name(&flag_name);

            script.push_str(&format!("complete -c {} -l {} ", cmd_name, flag_name));

            if let Some(desc) = &opt.description {
                // Escape single quotes in description
                let escaped_desc = desc.replace('\'', "'\\''");
                script.push_str(&format!("-d '{}' ", escaped_desc));
            }

            if opt.takes_argument {
                script.push_str("-r "); // require argument
            }

            script.push_str("\n");
        }

        // Process short flags
        for short_flag in &opt.short_flags {
            // Strip - prefix and clean trailing ...
            let mut flag_name = short_flag.trim_start_matches("-").to_string();
            flag_name = clean_flag_name(&flag_name);

            script.push_str(&format!("complete -c {} -s {} ", cmd_name, flag_name));

            if let Some(desc) = &opt.description {
                let escaped_desc = desc.replace('\'', "'\\''");
                script.push_str(&format!("-d '{}' ", escaped_desc));
            }

            if opt.takes_argument {
                script.push_str("-r ");
            }

            script.push_str("\n");
        }
    }

    // Add subcommands as a separate complete command
    // Use condition to show subcommands only when no flag is present
    let clean_subcommands = collect_clean_subcommands(result);
    if !clean_subcommands.is_empty() {
        if !clean_subcommands.is_empty() {
            // Condition: show subcommands when we're in a position to use a subcommand
            // -f flag prevents file completion (shows only the subcommands, not files)
            // __fish_use_subcommand is a built-in Fish function that returns true when
            // we're at a position where a subcommand can be used (not in a flag context)
            script.push_str(&format!(
                "complete -c {} -f -n '__fish_use_subcommand' -a '{}'\n",
                cmd_name,
                clean_subcommands.join(" ")
            ));
        }
    }

    script
}

/// Generate PowerShell completion script.
fn generate_powershell_completion(result: &ProbeResult) -> String {
    let cmd_name = &result.command;

    let mut script = String::new();
    script.push_str(&format!("# PowerShell completion for {}\n", cmd_name));
    script.push_str(&format!("# Generated by help-probe\n\n"));

    script.push_str(&format!(
        "Register-ArgumentCompleter -Native -CommandName '{}' -ScriptBlock {{\n",
        cmd_name
    ));
    script.push_str("    param($wordToComplete, $commandAst, $cursorPosition)\n\n");

    // Collect options (clean trailing ... from flags)
    let mut all_options: Vec<String> = Vec::new();
    for opt in &result.options {
        all_options.extend(opt.long_flags.iter().map(|f| clean_flag_name(f)));
    }

    // Collect subcommands (clean trailing commas)
    let subcommands = collect_clean_subcommands(result);

    script.push_str(&format!(
        "    $options = @({})\n",
        all_options
            .iter()
            .map(|o| format!("'{}'", o))
            .collect::<Vec<_>>()
            .join(", ")
    ));

    if !subcommands.is_empty() {
        script.push_str(&format!(
            "    $subcommands = @({})\n",
            subcommands
                .iter()
                .map(|s| format!("'{}'", s))
                .collect::<Vec<_>>()
                .join(", ")
        ));
    }

    script.push_str("\n");
    script.push_str("    if ($wordToComplete -match '^-') {\n");
    script.push_str("        $options | Where-Object { $_ -like \"$wordToComplete*\" }\n");
    script.push_str("    } else {\n");
    if !subcommands.is_empty() {
        script.push_str("        $subcommands | Where-Object { $_ -like \"$wordToComplete*\" }\n");
    } else {
        script.push_str(
            "        Get-ChildItem | Where-Object { $_.Name -like \"$wordToComplete*\" }\n",
        );
    }
    script.push_str("    }\n");
    script.push_str("}\n");

    script
}

/// Generate Nushell completion script.
///
/// Nushell uses the `extern` command to define external command signatures.
/// This generates a completion script that can be added to config.nu.
fn generate_nushell_completion(result: &ProbeResult) -> String {
    let cmd_name = &result.command;

    let mut script = String::new();
    script.push_str(&format!("# Nushell completion for {}\n", cmd_name));
    script.push_str(&format!("# Generated by help-probe\n"));
    script.push_str(&format!("# Add this to your config.nu or source it\n\n"));

    // Generate extern command signature
    script.push_str(&format!("extern {} [\n", cmd_name));

    // Add options as flags
    for opt in &result.options {
        for long_flag in &opt.long_flags {
            // Strip -- prefix and clean trailing ... (Nushell doesn't allow ... in flag names)
            let mut flag_name = long_flag.trim_start_matches("--").to_string();
            flag_name = clean_flag_name(&flag_name);

            // Nushell syntax: --flag: type (not --flag: argname: type)
            let flag_type = if opt.takes_argument {
                infer_nushell_type(opt).to_string()
            } else {
                "".to_string()
            };

            script.push_str(&format!("  --{}", flag_name));
            if !flag_type.is_empty() {
                script.push_str(&format!(": {}", flag_type));
            }
            if let Some(desc) = &opt.description {
                script.push_str(&format!("  # {}", desc.replace('\n', " ")));
            }
            script.push_str("\n");
        }

        // Add short flags
        for short_flag in &opt.short_flags {
            // Strip - prefix and clean trailing ... (Nushell doesn't allow ... in flag names)
            let mut flag_name = short_flag.trim_start_matches("-").to_string();
            flag_name = clean_flag_name(&flag_name);

            script.push_str(&format!("  -{}", flag_name));
            if opt.takes_argument {
                script.push_str(&format!(": {}", infer_nushell_type(opt)));
            }
            if let Some(desc) = &opt.description {
                script.push_str(&format!("  # {}", desc.replace('\n', " ")));
            }
            script.push_str("\n");
        }
    }

    // Add subcommands as positional arguments with completion
    if !result.subcommands.is_empty() {
        script.push_str("  subcommand?: string  # Subcommand to run\n");
    }

    // Add other arguments (only if no subcommands, or as additional args)
    // Nushell extern format: positional args come after flags
    if result.subcommands.is_empty() {
        for arg in &result.arguments {
            let arg_type = arg
                .arg_type
                .as_ref()
                .map(|t| match t {
                    crate::model::ArgumentType::Path => "path".to_string(),
                    crate::model::ArgumentType::Number => "number".to_string(),
                    crate::model::ArgumentType::Url => "string".to_string(),
                    crate::model::ArgumentType::Email => "string".to_string(),
                    _ => "string".to_string(),
                })
                .unwrap_or_else(|| "string".to_string());

            let marker = if arg.required { "" } else { "?" };
            let variadic = if arg.variadic { "..." } else { "" };
            let arg_name = arg.name.to_lowercase().replace(['<', '>'], "");
            script.push_str(&format!(
                "  {}{}: {}  # {}{}\n",
                arg_name,
                marker,
                arg_type,
                variadic,
                arg.description.as_deref().unwrap_or("argument")
            ));
        }
    } else {
        // If there are subcommands, add a catch-all for additional args
        if !result.arguments.is_empty() {
            script.push_str("  ...args: string  # Additional arguments\n");
        }
    }

    script.push_str("]\n\n");

    // Add custom completion for subcommands if they exist
    if !result.subcommands.is_empty() {
        let completer_name = format!("nu-complete-{}", cmd_name.replace('-', "_"));
        script.push_str(&format!("\n# Custom completion function for subcommands\n"));
        script.push_str(&format!("def {} [] {{\n", completer_name));
        script.push_str("  [\n");
        for subcmd in &result.subcommands {
            // Clean subcommand name: remove trailing commas and invalid entries
            let Some(clean_name) = clean_subcommand_name(&subcmd.name) else {
                continue;
            };

            script.push_str(&format!("    \"{}\"", clean_name));
            if let Some(desc) = &subcmd.description {
                // Clean up description for comment
                let clean_desc = desc.replace('\n', " ").trim().to_string();
                if !clean_desc.is_empty() {
                    script.push_str(&format!("  # {}", clean_desc));
                }
            }
            script.push_str("\n");
        }
        script.push_str("  ]\n");
        script.push_str("}\n\n");

        // Update extern to use the completer
        script.push_str(&format!(
            "# To enable subcommand completion, replace the subcommand line above with:\n"
        ));
        script.push_str(&format!(
            "#   subcommand?: string@{}  # Subcommand to run\n",
            completer_name
        ));
    }

    script
}

/// Infer Nushell type from option metadata.
fn infer_nushell_type(opt: &crate::model::OptionSpec) -> &'static str {
    match opt.option_type {
        crate::model::OptionType::Number => "number",
        crate::model::OptionType::Path => "path",
        crate::model::OptionType::Boolean => "bool",
        _ => "string",
    }
}