use std::io::Write;
use std::path::PathBuf;
use anyhow::{bail, Context, Result};
use clap::CommandFactory;
use clap_complete::{generate, Shell};
use crate::cli::args::{Cli, CompletionCommands};
const COMPLETION_MARKER: &str = "# zilliz-cli completion";
pub fn run(cmd: CompletionCommands) -> Result<()> {
match cmd {
CompletionCommands::Install { shell, apply } => install(shell, apply),
CompletionCommands::Uninstall { shell } => uninstall(shell),
CompletionCommands::Status { shell } => status(shell),
CompletionCommands::Show { shell } => show(shell),
}
}
fn detect_shell() -> Result<Shell> {
let shell_env = std::env::var("SHELL").unwrap_or_default();
if shell_env.contains("zsh") {
Ok(Shell::Zsh)
} else if shell_env.contains("bash") {
Ok(Shell::Bash)
} else if shell_env.contains("fish") {
Ok(Shell::Fish)
} else {
bail!("Cannot detect shell. Specify one: bash, zsh, or fish")
}
}
fn resolve_shell(shell: Option<Shell>) -> Result<Shell> {
match shell {
Some(s) => Ok(s),
None => detect_shell(),
}
}
fn rc_path(shell: Shell) -> PathBuf {
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("~"));
match shell {
Shell::Bash => home.join(".bashrc"),
Shell::Zsh => home.join(".zshrc"),
Shell::Fish => home.join(".config/fish/config.fish"),
_ => home.join(".bashrc"),
}
}
fn completion_script_line(shell: Shell) -> String {
match shell {
Shell::Bash => r#"eval "$(zilliz completion show bash)""#.to_string(),
Shell::Zsh => r#"eval "$(zilliz completion show zsh)""#.to_string(),
Shell::Fish => "zilliz completion show fish | source".to_string(),
_ => r#"eval "$(zilliz completion show bash)""#.to_string(),
}
}
fn is_installed(shell: Shell) -> bool {
let path = rc_path(shell);
if !path.exists() {
return false;
}
let content = std::fs::read_to_string(&path).unwrap_or_default();
content.contains(COMPLETION_MARKER) || content.contains("zilliz completion show")
}
fn install(shell: Option<Shell>, apply: bool) -> Result<()> {
let shell = resolve_shell(shell)?;
if is_installed(shell) {
println!("Completion already installed for {:?}.", shell);
return Ok(());
}
let path = rc_path(shell);
let script = completion_script_line(shell);
if apply {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut f = std::fs::OpenOptions::new()
.append(true)
.create(true)
.open(&path)
.with_context(|| format!("Cannot write to {}", path.display()))?;
writeln!(f, "\n{}", COMPLETION_MARKER)?;
writeln!(f, "{}", script)?;
println!("Completion installed to {}", path.display());
println!("Restart your shell or run: source {}", path.display());
} else {
println!("Add this to {}:", path.display());
println!();
println!(" {}", script);
println!();
println!("Or run: zilliz completion install {:?} --apply", shell);
}
Ok(())
}
fn uninstall(shell: Option<Shell>) -> Result<()> {
let shell = resolve_shell(shell)?;
let path = rc_path(shell);
if !path.exists() {
println!("RC file not found: {}", path.display());
return Ok(());
}
let content = std::fs::read_to_string(&path)?;
let lines: Vec<&str> = content.lines().collect();
let mut new_lines = Vec::new();
let mut skip_next = false;
for line in &lines {
if skip_next {
skip_next = false;
continue;
}
if line.contains(COMPLETION_MARKER) {
skip_next = true; continue;
}
if line.contains("zilliz completion show") {
continue;
}
new_lines.push(*line);
}
std::fs::write(&path, new_lines.join("\n") + "\n")?;
println!("Completion removed from {}", path.display());
Ok(())
}
fn status(shell: Option<Shell>) -> Result<()> {
let shell = resolve_shell(shell)?;
if is_installed(shell) {
println!("Completion is installed for {:?}.", shell);
} else {
println!("Completion is NOT installed for {:?}.", shell);
println!("Run: zilliz completion install {:?} --apply", shell);
}
Ok(())
}
fn show(shell: Shell) -> Result<()> {
let mut cli = Cli::command();
generate(shell, &mut cli, "zilliz", &mut std::io::stdout());
Ok(())
}