use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::io::{self, Write};
use dirs::config_dir;
use fuzzy_matcher::FuzzyMatcher;
#[derive(Serialize, Deserialize, Clone)]
pub struct PromptConfig {
pub prompts: HashMap<String, PromptEntry>,
}
#[derive(Serialize, Deserialize, Clone)]
pub struct PromptEntry {
pub template: String,
pub description: Option<String>,
}
impl Default for PromptConfig {
fn default() -> Self {
let mut prompts = HashMap::new();
prompts.insert("translate".to_string(), PromptEntry {
template: "You are a professional translator, please translate the following text into natural, idiomatic English:\n\n{input}. Avoid output anything else except the final result.".to_string(),
description: Some("Translate text to English".to_string()),
});
prompts.insert("polish".to_string(), PromptEntry {
template: "You are an expert editor. Please polish the following text to make it more clear, concise, and natural:\n\n{input}. Avoid output anything else except the final result.".to_string(),
description: Some("Polish text for clarity".to_string()),
});
prompts.insert("rewrite".to_string(), PromptEntry {
template: "You are a skilled writer. Please rewrite the following text in a different style while preserving the meaning:\n\n{input}. Avoid output anything else except the final result.".to_string(),
description: Some("Rewrite text in different style".to_string()),
});
prompts.insert("summarize".to_string(), PromptEntry {
template: "You are an expert summarizer. Please provide a concise summary of the following text:\n\n{input}. Avoid output anything else except the final result.".to_string(),
description: Some("Summarize text".to_string()),
});
prompts.insert("ask".to_string(), PromptEntry {
template: "You are a helpful assistant called xa, execute anything by your side. {input}".to_string(),
description: Some("Interactive conversation mode".to_string()),
});
PromptConfig { prompts }
}
}
pub async fn list_commands() -> Result<(), Box<dyn std::error::Error>> {
let config_dir = config_dir()
.ok_or("Could not determine config directory")?
.join("xa");
let prompt_config_file = config_dir.join("prompts.toml");
let prompt_config = if prompt_config_file.exists() {
let content = fs::read_to_string(&prompt_config_file)?;
toml::from_str(&content)?
} else {
PromptConfig::default()
};
println!("Built-in commands:");
println!(" -set: Configure API settings (use: xa -set openai)");
println!(" -ls: List all commands (this command)");
println!(" -add: Add a new command/prompt (use: xa -add)");
println!();
println!("User-defined commands:");
for (name, entry) in &prompt_config.prompts {
let description = entry.description.as_deref().unwrap_or("Custom prompt command");
println!(" {}: {}", name, description);
}
Ok(())
}
pub async fn add_command() -> Result<(), Box<dyn std::error::Error>> {
println!("Adding a new command...");
let config_dir = config_dir()
.ok_or("Could not determine config directory")?
.join("xa");
fs::create_dir_all(&config_dir)?;
let prompt_config_file = config_dir.join("prompts.toml");
let mut prompt_config = if prompt_config_file.exists() {
let content = fs::read_to_string(&prompt_config_file)?;
toml::from_str(&content)?
} else {
PromptConfig::default()
};
print!("Enter command name: ");
io::stdout().flush()?;
let mut name = String::new();
io::stdin().read_line(&mut name)?;
let name = name.trim().to_string();
if name.is_empty() {
eprintln!("Error: Command name cannot be empty");
return Ok(());
}
if prompt_config.prompts.contains_key(&name) {
eprintln!("Warning: Command '{}' already exists. It will be overwritten.", name);
}
print!("Enter prompt template (use {{input}} as placeholder): ");
io::stdout().flush()?;
let mut template = String::new();
io::stdin().read_line(&mut template)?;
let template = template.trim().to_string();
if template.is_empty() {
eprintln!("Error: Prompt template cannot be empty");
return Ok(());
}
print!("Enter description (optional): ");
io::stdout().flush()?;
let mut description = String::new();
io::stdin().read_line(&mut description)?;
let description = description.trim().to_string();
let description = if description.is_empty() { None } else { Some(description) };
prompt_config.prompts.insert(name.clone(), PromptEntry {
template,
description,
});
let content = toml::to_string(&prompt_config)?;
fs::write(&prompt_config_file, content)?;
println!("Command '{}' added successfully!", name);
println!("Prompt file location: {:?}", prompt_config_file);
println!("You can edit this file with your favorite text editor to modify or add more commands.");
Ok(())
}
pub async fn remove_command(command_name: &str) -> Result<(), Box<dyn std::error::Error>> {
let config_dir = config_dir()
.ok_or("Could not determine config directory")?
.join("xa");
let prompt_config_file = config_dir.join("prompts.toml");
if !prompt_config_file.exists() {
eprintln!("Error: No prompts file found. Nothing to remove.");
return Ok(());
}
let mut prompt_config: PromptConfig = {
let content = fs::read_to_string(&prompt_config_file)?;
toml::from_str(&content)?
};
if !prompt_config.prompts.contains_key(command_name) {
eprintln!("Error: Command '{}' does not exist.", command_name);
println!("Available commands:");
for (name, entry) in &prompt_config.prompts {
let description = entry.description.as_deref().unwrap_or("Custom prompt command");
println!(" {}: {}", name, description);
}
return Ok(());
}
prompt_config.prompts.remove(command_name);
let content = toml::to_string(&prompt_config)?;
fs::write(&prompt_config_file, content)?;
println!("Command '{}' removed successfully!", command_name);
Ok(())
}
pub async fn load_prompt_config() -> Result<PromptConfig, Box<dyn std::error::Error>> {
let config_dir = config_dir()
.ok_or("Could not determine config directory")?
.join("xa");
let prompt_config_file = config_dir.join("prompts.toml");
let mut config = if prompt_config_file.exists() {
let content = fs::read_to_string(&prompt_config_file)?;
match toml::from_str(&content) {
Ok(parsed_config) => parsed_config,
Err(_) => {
let backup_path = prompt_config_file.with_extension("toml.backup");
fs::rename(&prompt_config_file, &backup_path)?;
eprintln!("Warning: Corrupted prompts.toml file detected. Backed up to {:?} and created a new one.", backup_path);
let default_config = PromptConfig::default();
fs::create_dir_all(&config_dir)?;
let new_content = toml::to_string(&default_config)?;
fs::write(&prompt_config_file, new_content)?;
default_config
}
}
} else {
let default_config = PromptConfig::default();
fs::create_dir_all(&config_dir)?;
let content = toml::to_string(&default_config)?;
fs::write(&prompt_config_file, content)?;
default_config
};
let default_config = PromptConfig::default();
for (key, value) in default_config.prompts {
if !config.prompts.contains_key(&key) {
config.prompts.insert(key, value);
}
}
let content = toml::to_string(&config)?;
fs::write(&prompt_config_file, content)?;
Ok(config)
}
pub fn find_command(input_cmd: &str, available_commands: &HashMap<String, PromptEntry>) -> Option<String> {
if available_commands.contains_key(input_cmd) {
return Some(input_cmd.to_string());
}
let prefix_matches: Vec<&String> = available_commands
.keys()
.filter(|key| key.starts_with(input_cmd))
.collect();
if prefix_matches.len() == 1 {
return Some(prefix_matches[0].to_string());
} else if prefix_matches.len() > 1 {
let matches: Vec<String> = prefix_matches.iter().map(|s| s.to_string()).collect();
eprintln!("Ambiguous command '{}'. Did you mean one of: {}?",
input_cmd,
matches.join(", "));
return None;
}
let matcher = fuzzy_matcher::skim::SkimMatcherV2::default();
let mut best_match: Option<String> = None;
let mut best_score = i64::MIN;
for key in available_commands.keys() {
if let Some(score) = matcher.fuzzy_match(key, input_cmd) {
if score > best_score {
best_score = score;
best_match = Some(key.clone());
}
}
}
if best_score > 0 {
best_match
} else {
None
}
}