use crate::api_key::{ApiKeyProvider, SecureString};
use crate::error::{Error, Result};
use crate::gemini::GeminiClient;
use crate::quota::QuotaTracker;
use colored::*;
use dialoguer::{theme::ColorfulTheme, Confirm, Password};
use std::sync::Arc;
const GOOGLE_AI_STUDIO_URL: &str = "https://aistudio.google.com/apikey";
pub struct SetupWizard {
api_key_provider: Option<Arc<dyn ApiKeyProvider>>,
}
impl SetupWizard {
pub fn new() -> Self {
Self {
api_key_provider: None,
}
}
pub fn with_provider(provider: Arc<dyn ApiKeyProvider>) -> Self {
Self {
api_key_provider: Some(provider),
}
}
pub async fn run(&self) -> Result<SecureString> {
self.print_welcome();
let should_open = Confirm::with_theme(&ColorfulTheme::default())
.with_prompt("Would you like to open Google AI Studio in your browser?")
.default(true)
.interact()
.map_err(|e| Error::ConfigError(format!("Failed to get user input: {}", e)))?;
if should_open {
self.open_api_key_page()?;
} else {
println!(
"\n{} {}",
"→".cyan().bold(),
format!("Visit: {}", GOOGLE_AI_STUDIO_URL).underline()
);
}
self.print_instructions();
let api_key = self.prompt_for_api_key()?;
self.validate_and_test_key(&api_key).await?;
if let Some(provider) = &self.api_key_provider {
provider
.set_key(api_key.clone())
.await
.map_err(|e| Error::ConfigError(format!("Failed to store API key: {}", e)))?;
println!(
"\n{} {}",
"✓".green().bold(),
"API key securely stored!".green()
);
}
self.print_completion();
Ok(api_key)
}
fn print_welcome(&self) {
println!("\n{}", "━".repeat(70).cyan());
println!(
"{} {}",
"🚀".to_string().bold(),
"Welcome to Dynamic Grounding Setup".bold().cyan()
);
println!("{}", "━".repeat(70).cyan());
println!(
"\n{}",
"This wizard will help you set up Google Gemini AI integration.".white()
);
println!(
"{}",
"You'll need a free API key from Google AI Studio.".white()
);
}
fn open_api_key_page(&self) -> Result<()> {
println!(
"\n{} {}",
"→".cyan().bold(),
"Opening Google AI Studio in your browser...".cyan()
);
match open::that(GOOGLE_AI_STUDIO_URL) {
Ok(_) => {
println!(
"{} {}",
"✓".green().bold(),
"Browser opened successfully!".green()
);
Ok(())
}
Err(e) => {
println!(
"{} {}",
"⚠".yellow().bold(),
format!("Could not open browser automatically: {}", e).yellow()
);
println!(
"{} {}",
"→".cyan().bold(),
format!("Please visit: {}", GOOGLE_AI_STUDIO_URL)
.underline()
.cyan()
);
Ok(())
}
}
}
fn print_instructions(&self) {
println!("\n{}", "Instructions:".bold().white());
println!(" {} Sign in with your Google account", "1.".cyan().bold());
println!(
" {} Click '{}' or '{}' button",
"2.".cyan().bold(),
"Create API key".yellow(),
"Get API key".yellow()
);
println!(" {} Copy the generated API key", "3.".cyan().bold());
println!(
" {} Paste it below (it won't be displayed)\n",
"4.".cyan().bold()
);
println!("{}", "Free Tier Limits:".bold().white());
println!(" • {} requests per minute", "15".green().bold());
println!(" • {} requests per day", "1,500".green().bold());
println!(" • {} tokens per minute", "1,000,000".green().bold());
println!(" • {} token context window\n", "2,000,000".green().bold());
}
fn prompt_for_api_key(&self) -> Result<SecureString> {
loop {
let key = Password::with_theme(&ColorfulTheme::default())
.with_prompt("Enter your Google Gemini API key")
.interact()
.map_err(|e| Error::ConfigError(format!("Failed to read API key: {}", e)))?;
if key.is_empty() {
println!("{} API key cannot be empty\n", "✗".red().bold());
continue;
}
if !key.starts_with("AIza") || key.len() < 30 {
println!(
"{} {}",
"⚠".yellow().bold(),
"This doesn't look like a valid Google API key".yellow()
);
println!(
" {}",
"Google API keys start with 'AIza' and are typically 39 characters long"
.dimmed()
);
let should_continue = Confirm::with_theme(&ColorfulTheme::default())
.with_prompt("Do you want to try again?")
.default(true)
.interact()
.map_err(|e| Error::ConfigError(format!("Failed to get user input: {}", e)))?;
if !should_continue {
return Err(Error::ApiKeyError("Invalid API key format".to_string()));
}
continue;
}
return Ok(SecureString::new(key));
}
}
async fn validate_and_test_key(&self, api_key: &SecureString) -> Result<()> {
println!("\n{} {}", "→".cyan().bold(), "Testing API key...".cyan());
struct TestProvider {
key: SecureString,
}
#[async_trait::async_trait]
impl ApiKeyProvider for TestProvider {
async fn get_key(&self) -> Result<SecureString> {
Ok(self.key.clone())
}
}
let provider = Arc::new(TestProvider {
key: api_key.clone(),
});
let quota_tracker = Arc::new(QuotaTracker::new());
let client = GeminiClient::new(provider, quota_tracker);
match client
.generate_content("Respond with just 'OK' if you receive this message.", None)
.await
{
Ok(_response) => {
println!(
"{} {}",
"✓".green().bold(),
"API key is valid and working!".green()
);
Ok(())
}
Err(e) => {
println!(
"{} {}",
"✗".red().bold(),
format!("API key validation failed: {}", e).red()
);
println!("\n{}", "Common issues:".bold().yellow());
println!(" • Make sure you copied the entire key");
println!(" • Check that the API key is enabled in Google AI Studio");
println!(" • Verify you have access to the Gemini API");
Err(Error::ApiKeyError(format!(
"API key validation failed: {}",
e
)))
}
}
}
fn print_completion(&self) {
println!("\n{}", "━".repeat(70).cyan());
println!(
"{} {}",
"✓".green().bold(),
"Setup Complete!".bold().green()
);
println!("{}", "━".repeat(70).cyan());
println!(
"\n{} You're ready to use Dynamic Grounding with GitHub Copilot!",
"🎉".to_string().bold()
);
println!("\n{}", "Next steps:".bold().white());
println!(" • Your API key is securely stored");
println!(" • The MCP server will start automatically");
println!(" • GitHub Copilot can now use Gemini for enhanced code search\n");
}
pub async fn quick_setup(
api_key: String,
provider: Option<Arc<dyn ApiKeyProvider>>,
) -> Result<()> {
let secure_key = SecureString::new(api_key);
struct TestProvider {
key: SecureString,
}
#[async_trait::async_trait]
impl ApiKeyProvider for TestProvider {
async fn get_key(&self) -> Result<SecureString> {
Ok(self.key.clone())
}
}
let test_provider = Arc::new(TestProvider {
key: secure_key.clone(),
});
let quota_tracker = Arc::new(QuotaTracker::new());
let client = GeminiClient::new(test_provider, quota_tracker);
client
.generate_content("Respond with just 'OK' if you receive this message.", None)
.await?;
if let Some(p) = provider {
p.set_key(secure_key).await?;
}
Ok(())
}
}
impl Default for SetupWizard {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_wizard_creation() {
let wizard = SetupWizard::new();
assert!(wizard.api_key_provider.is_none());
let wizard = SetupWizard::default();
assert!(wizard.api_key_provider.is_none());
}
#[test]
fn test_api_key_url() {
assert_eq!(GOOGLE_AI_STUDIO_URL, "https://aistudio.google.com/apikey");
}
}