Skip to content
use std::fmt;
use std::fmt::Write as _;
use std::str::FromStr;

use crate::nanoargs::parser::ArgParser;

/// Supported shell targets for completion script generation.
///
/// Use with [`ArgParser::generate_completions`] to produce a shell-specific
/// completion script from your CLI schema. The generated script can be sourced
/// directly or installed into the shell's completion directory.
///
/// `Shell` implements [`FromStr`] (case-insensitive) and [`Display`], so it
/// works naturally as a CLI argument value:
///
/// ```
/// use nanoargs::Shell;
///
/// let shell: Shell = "zsh".parse().unwrap();
/// assert_eq!(shell, Shell::Zsh);
/// assert_eq!(shell.to_string(), "zsh");
/// ```
///
/// Accepted string values: `bash`, `zsh`, `fish`, `powershell` (or `pwsh`).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[allow(clippy::enum_variant_names)]
pub enum Shell {
    /// GNU Bourne-Again Shell.
    Bash,
    /// Z Shell — supports descriptions and mutual-exclusion groups.
    Zsh,
    /// Friendly Interactive Shell.
    Fish,
    /// PowerShell (cross-platform).
    PowerShell,
}

impl FromStr for Shell {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "bash" => Ok(Shell::Bash),
            "zsh" => Ok(Shell::Zsh),
            "fish" => Ok(Shell::Fish),
            "powershell" | "pwsh" => Ok(Shell::PowerShell),
            _ => Err(format!("unsupported shell: {s}")),
        }
    }
}

impl fmt::Display for Shell {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Shell::Bash => write!(f, "bash"),
            Shell::Zsh => write!(f, "zsh"),
            Shell::Fish => write!(f, "fish"),
            Shell::PowerShell => write!(f, "powershell"),
        }
    }
}

impl ArgParser {
    /// Generate a shell completion script for the given shell.
    ///
    /// Returns a self-contained script string that, when sourced or installed,
    /// provides tab-completion for your CLI's flags, options, and subcommands.
    ///
    /// # Shell-specific behaviour
    ///
    /// | Shell | Descriptions | Short/long exclusion | Script format |
    /// |-------|-------------|---------------------|---------------|
    /// | Bash | No | No | `complete -F` function |
    /// | Zsh | Yes | Yes | `#compdef` + `_arguments` |
    /// | Fish | Yes | No | `complete -c` commands |
    /// | PowerShell | Yes | No | `Register-ArgumentCompleter` |
    ///
    /// # Example
    ///
    /// ```
    /// use nanoargs::{ArgBuilder, Flag, Shell};
    ///
    /// let parser = ArgBuilder::new()
    ///     .name("myapp")
    ///     .flag(Flag::new("verbose").short('v').desc("Enable verbose output"))
    ///     .build()
    ///     .unwrap();
    ///
    /// let script = parser.generate_completions(Shell::Bash);
    /// assert!(script.contains("complete -F"));
    /// ```
    pub fn generate_completions(&self, shell: Shell) -> String {
        match shell {
            Shell::Bash => generate_bash(self),
            Shell::Zsh => generate_zsh(self),
            Shell::Fish => generate_fish(self),
            Shell::PowerShell => generate_powershell(self),
        }
    }
}

fn generate_bash(parser: &ArgParser) -> String {
    let name = parser.program_name().unwrap_or("program");
    let func_name = name.replace('-', "_");
    let mut out = String::new();

    writeln!(out, "_{func_name}() {{").unwrap();
    out.push_str("    local i cur prev opts cmd\n");
    out.push_str("    COMPREPLY=()\n");
    out.push_str("    if type _init_completion &>/dev/null; then\n");
    out.push_str("        _init_completion || return\n");
    out.push_str("    else\n");
    out.push_str("        cur=\"${COMP_WORDS[COMP_CWORD]}\"\n");
    out.push_str("        prev=\"${COMP_WORDS[COMP_CWORD-1]}\"\n");
    out.push_str("    fi\n\n");

    let subs = parser.subcommands();
    if subs.is_empty() {
        // No subcommands — flat completion
        let words = bash_completable_words(parser);
        writeln!(out, "    opts=\"{words}\"").unwrap();
        out.push_str("    COMPREPLY=($(compgen -W \"${opts}\" -- \"${cur}\"))\n");
    } else {
        // With subcommands — track which subcommand we're in via a cmd variable
        out.push_str("    cmd=\"\"\n");
        out.push_str("    for i in \"${COMP_WORDS[@]}\"; do\n");
        out.push_str("        case \"${i}\" in\n");
        for sub in subs {
            writeln!(out, "            {}) cmd=\"{}\";;", sub.name, sub.name).unwrap();
        }
        out.push_str("        esac\n");
        out.push_str("    done\n\n");

        out.push_str("    case \"${cmd}\" in\n");
        for sub in subs {
            let sub_words = bash_completable_words(&sub.parser);
            writeln!(out, "        {})", sub.name).unwrap();
            writeln!(out, "            opts=\"{sub_words}\"").unwrap();
            out.push_str("            ;;\n");
        }
        // Default: top-level options + subcommand names
        let sub_names: Vec<&str> = subs.iter().map(|s| s.name.as_str()).collect();
        let top_words = bash_completable_words(parser);
        let all_top = format!("{} {}", sub_names.join(" "), top_words);
        out.push_str("        *)\n");
        writeln!(out, "            opts=\"{}\"", all_top.trim()).unwrap();
        out.push_str("            ;;\n");
        out.push_str("    esac\n\n");

        out.push_str("    COMPREPLY=($(compgen -W \"${opts}\" -- \"${cur}\"))\n");
    }

    out.push_str("}\n\n");
    writeln!(
        out,
        "complete -F _{func_name} -o bashdefault -o default {name}"
    )
    .unwrap();
    out
}

/// Escape single quotes for embedding in bash single-quoted strings.
fn bash_escape(s: &str) -> String {
    s.replace('\'', "'\\''")
}

/// Build the completable word list for a parser (flags + options, non-hidden).
fn bash_completable_words(parser: &ArgParser) -> String {
    let mut words = Vec::new();
    for flag in parser.flags() {
        if flag.hidden {
            continue;
        }
        words.push(format!("--{}", flag.long));
        if let Some(c) = flag.short {
            words.push(format!("-{c}"));
        }
    }
    for opt in parser.options() {
        if opt.hidden {
            continue;
        }
        words.push(format!("--{}", opt.long));
        if let Some(c) = opt.short {
            words.push(format!("-{c}"));
        }
    }
    bash_escape(&words.join(" "))
}

fn generate_zsh(parser: &ArgParser) -> String {
    let name = parser.program_name().unwrap_or("program");
    let func_name = name.replace('-', "_");
    let mut out = String::new();

    write!(out, "#compdef {name}\n\n").unwrap();

    let subs = parser.subcommands();
    if subs.is_empty() {
        // Simple case: no subcommands
        writeln!(out, "_{func_name}() {{").unwrap();
        out.push_str("    _arguments \\\n");
        zsh_push_args(parser, &mut out, "        ");
        out.push_str("}\n\n");
    } else {
        // With subcommands: main function dispatches
        writeln!(out, "_{func_name}() {{").unwrap();
        out.push_str("    local line state\n\n");

        // Top-level flags/options
        out.push_str("    _arguments \\\n");
        zsh_push_args(parser, &mut out, "        ");

        // Subcommand dispatch
        out.push_str("        '*::subcmd:->subcmd' && return\n\n");
        out.push_str("    case $state in\n");
        out.push_str("        subcmd)\n");
        out.push_str("            case $line[1] in\n");
        for sub in subs {
            let sub_func = format!("_{func_name}_{}", sub.name.replace('-', "_"));
            writeln!(out, "                {})", sub.name).unwrap();
            writeln!(out, "                    {sub_func}").unwrap();
            out.push_str("                    ;;\n");
        }

        // If no subcommand matched yet, describe available subcommands
        out.push_str("                *)\n");
        out.push_str("                    local -a subcmds\n");
        out.push_str("                    subcmds=(\n");
        for sub in subs {
            let desc = zsh_escape(&sub.description);
            writeln!(out, "                        '{}:{}'", sub.name, desc).unwrap();
        }
        out.push_str("                    )\n");
        out.push_str("                    _describe 'subcommand' subcmds\n");
        out.push_str("                    ;;\n");
        out.push_str("            esac\n");
        out.push_str("            ;;\n");
        out.push_str("    esac\n");
        out.push_str("}\n\n");

        // Generate sub-functions for each subcommand
        for sub in subs {
            let sub_func = format!("_{func_name}_{}", sub.name.replace('-', "_"));
            writeln!(out, "{sub_func}() {{").unwrap();
            out.push_str("    _arguments \\\n");
            zsh_push_args(&sub.parser, &mut out, "        ");
            out.push_str("}\n\n");
        }
    }

    // Register the completion function with Zsh's completion system.
    // Works both when sourced directly and when loaded via fpath
    // (the #compdef header handles the fpath case).
    writeln!(out, "compdef _{func_name} {name}").unwrap();
    out
}

/// Escape brackets and colons for Zsh _arguments specs.
fn zsh_escape(s: &str) -> String {
    s.replace('[', "\\[")
        .replace(']', "\\]")
        .replace(':', "\\:")
}

/// Push _arguments specs for all non-hidden flags and options.
fn zsh_push_args(parser: &ArgParser, out: &mut String, indent: &str) {
    for flag in parser.flags() {
        if flag.hidden {
            continue;
        }
        let desc = zsh_escape(&flag.description);
        if let Some(c) = flag.short {
            // Mutual exclusion group: using -r excludes --release and vice versa
            let excl = format!("-{c} --{}", flag.long);
            writeln!(out, "{indent}'({excl})--{}[{desc}]' \\", flag.long).unwrap();
            writeln!(out, "{indent}'({excl})-{c}[{desc}]' \\").unwrap();
        } else {
            writeln!(out, "{indent}'--{}[{desc}]' \\", flag.long).unwrap();
        }
    }
    for opt in parser.options() {
        if opt.hidden {
            continue;
        }
        let desc = zsh_escape(&opt.description);
        let ph = if opt.placeholder.is_empty() {
            "value".to_string()
        } else {
            opt.placeholder.clone()
        };
        if let Some(c) = opt.short {
            let excl = format!("-{c} --{}", opt.long);
            writeln!(out, "{indent}'({excl})--{}[{desc}]:{ph}:' \\", opt.long).unwrap();
            writeln!(out, "{indent}'({excl})-{c}[{desc}]:{ph}:' \\").unwrap();
        } else {
            writeln!(out, "{indent}'--{}[{desc}]:{ph}:' \\", opt.long).unwrap();
        }
    }
}

fn generate_fish(parser: &ArgParser) -> String {
    let name = parser.program_name().unwrap_or("program");
    let mut out = String::new();

    let subs = parser.subcommands();
    let has_subs = !subs.is_empty();

    // Top-level flags and options
    for flag in parser.flags() {
        if flag.hidden {
            continue;
        }
        let mut cmd = format!("complete -c {name}");
        if has_subs {
            cmd.push_str(" -n '__fish_use_subcommand'");
        }
        write!(cmd, " -l {}", flag.long).unwrap();
        if let Some(c) = flag.short {
            write!(cmd, " -s {c}").unwrap();
        }
        if !flag.description.is_empty() {
            write!(cmd, " -d '{}'", fish_escape(&flag.description)).unwrap();
        }
        out.push_str(&cmd);
        out.push('\n');
    }

    for opt in parser.options() {
        if opt.hidden {
            continue;
        }
        let mut cmd = format!("complete -c {name}");
        if has_subs {
            cmd.push_str(" -n '__fish_use_subcommand'");
        }
        write!(cmd, " -l {}", opt.long).unwrap();
        if let Some(c) = opt.short {
            write!(cmd, " -s {c}").unwrap();
        }
        cmd.push_str(" -r");
        if !opt.description.is_empty() {
            write!(cmd, " -d '{}'", fish_escape(&opt.description)).unwrap();
        }
        out.push_str(&cmd);
        out.push('\n');
    }

    // Subcommands
    for sub in subs {
        // Register the subcommand name itself
        let mut cmd = format!(
            "complete -c {name} -n '__fish_use_subcommand' -a {}",
            sub.name
        );
        if !sub.description.is_empty() {
            write!(cmd, " -d '{}'", fish_escape(&sub.description)).unwrap();
        }
        out.push_str(&cmd);
        out.push('\n');

        // Subcommand-specific flags
        for flag in sub.parser.flags() {
            if flag.hidden {
                continue;
            }
            let mut cmd = format!(
                "complete -c {name} -n '__fish_seen_subcommand_from {}' -l {}",
                sub.name, flag.long
            );
            if let Some(c) = flag.short {
                write!(cmd, " -s {c}").unwrap();
            }
            if !flag.description.is_empty() {
                write!(cmd, " -d '{}'", fish_escape(&flag.description)).unwrap();
            }
            out.push_str(&cmd);
            out.push('\n');
        }

        // Subcommand-specific options
        for opt in sub.parser.options() {
            if opt.hidden {
                continue;
            }
            let mut cmd = format!(
                "complete -c {name} -n '__fish_seen_subcommand_from {}' -l {}",
                sub.name, opt.long
            );
            if let Some(c) = opt.short {
                write!(cmd, " -s {c}").unwrap();
            }
            cmd.push_str(" -r");
            if !opt.description.is_empty() {
                write!(cmd, " -d '{}'", fish_escape(&opt.description)).unwrap();
            }
            out.push_str(&cmd);
            out.push('\n');
        }
    }

    out
}

/// Escape single quotes for embedding in fish single-quoted strings.
fn fish_escape(s: &str) -> String {
    s.replace('\'', "'\\''")
}

fn generate_powershell(parser: &ArgParser) -> String {
    let name = parser.program_name().unwrap_or("program");
    let mut out = String::new();

    writeln!(
        out,
        "Register-ArgumentCompleter -CommandName '{name}' -ScriptBlock {{"
    )
    .unwrap();
    out.push_str("    param($commandName, $wordToComplete, $cursorPosition)\n\n");

    let subs = parser.subcommands();
    if subs.is_empty() {
        // No subcommands — flat completion
        out.push_str("    $completions = @(\n");
        ps_push_completions(parser, &mut out, "        ");
        out.push_str("    )\n\n");
        out.push_str(
            "    $completions | Where-Object { $_.CompletionText -like \"$wordToComplete*\" }\n",
        );
    } else {
        // Find the current subcommand
        out.push_str("    # Determine the subcommand\n");
        out.push_str("    $tokens = $commandName -split '\\s+'\n");
        out.push_str("    $subcmd = $null\n");
        out.push_str("    foreach ($t in $tokens[1..($tokens.Length - 1)]) {\n");
        out.push_str("        switch ($t) {\n");
        for sub in subs {
            writeln!(
                out,
                "            '{}' {{ $subcmd = '{}' }}",
                sub.name, sub.name
            )
            .unwrap();
        }
        out.push_str("        }\n");
        out.push_str("    }\n\n");

        out.push_str("    $completions = @()\n\n");
        out.push_str("    switch ($subcmd) {\n");
        for sub in subs {
            writeln!(out, "        '{}' {{", sub.name).unwrap();
            out.push_str("            $completions = @(\n");
            ps_push_completions(&sub.parser, &mut out, "                ");
            out.push_str("            )\n");
            out.push_str("        }\n");
        }
        // Default: top-level completions + subcommand names
        out.push_str("        default {\n");
        out.push_str("            $completions = @(\n");
        ps_push_completions(parser, &mut out, "                ");
        for sub in subs {
            let desc = ps_escape(&sub.description);
            writeln!(out,
                "                [System.Management.Automation.CompletionResult]::new('{}', '{}', 'ParameterValue', '{}')",
                sub.name, sub.name, desc
            ).unwrap();
        }
        out.push_str("            )\n");
        out.push_str("        }\n");
        out.push_str("    }\n\n");

        out.push_str(
            "    $completions | Where-Object { $_.CompletionText -like \"$wordToComplete*\" }\n",
        );
    }

    out.push_str("}\n");
    out
}
/// Escape special PowerShell characters in description strings.
fn ps_escape(s: &str) -> String {
    s.replace('\'', "''")
}

/// Push [`CompletionResult`]`::new()` entries for all non-hidden flags and options.
fn ps_push_completions(parser: &ArgParser, out: &mut String, indent: &str) {
    for flag in parser.flags() {
        if flag.hidden {
            continue;
        }
        let desc = ps_escape(&flag.description);
        writeln!(out,
            "{indent}[System.Management.Automation.CompletionResult]::new('--{}', '--{}', 'ParameterName', '{}')",
            flag.long, flag.long, desc
        ).unwrap();
        if let Some(c) = flag.short {
            writeln!(out,
                "{indent}[System.Management.Automation.CompletionResult]::new('-{c}', '-{c}', 'ParameterName', '{desc}')"
            ).unwrap();
        }
    }
    for opt in parser.options() {
        if opt.hidden {
            continue;
        }
        let desc = ps_escape(&opt.description);
        writeln!(out,
            "{indent}[System.Management.Automation.CompletionResult]::new('--{}', '--{}', 'ParameterName', '{}')",
            opt.long, opt.long, desc
        ).unwrap();
        if let Some(c) = opt.short {
            writeln!(out,
                "{indent}[System.Management.Automation.CompletionResult]::new('-{c}', '-{c}', 'ParameterName', '{desc}')"
            ).unwrap();
        }
    }
}