dynamic_grounding_for_github_copilot 0.1.0

MCP server providing Google Gemini AI integration for enhanced codebase search and analysis
Documentation
//! Interactive setup wizard for first-time configuration.
//!
//! This module provides a user-friendly CLI wizard that:
//! - Opens the Google AI Studio API key page in the browser
//! - Provides clear instructions for obtaining an API key
//! - Validates the API key format
//! - Securely stores the key
//! - Tests the connection to verify the key works

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";

/// Interactive setup wizard for first-time configuration
pub struct SetupWizard {
    api_key_provider: Option<Arc<dyn ApiKeyProvider>>,
}

impl SetupWizard {
    /// Create a new setup wizard
    pub fn new() -> Self {
        Self {
            api_key_provider: None,
        }
    }

    /// Create a setup wizard with an existing API key provider
    pub fn with_provider(provider: Arc<dyn ApiKeyProvider>) -> Self {
        Self {
            api_key_provider: Some(provider),
        }
    }

    /// Run the interactive setup wizard
    pub async fn run(&self) -> Result<SecureString> {
        self.print_welcome();

        // Check if user wants to open the browser
        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();

        // Get the API key from user
        let api_key = self.prompt_for_api_key()?;

        // Validate the API key
        self.validate_and_test_key(&api_key).await?;

        // Store the key if we have a provider
        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)
    }

    /// Print welcome message
    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()
        );
    }

    /// Open the Google AI Studio API key page in the browser
    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(())
            }
        }
    }

    /// Print instructions for obtaining an API key
    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());
    }

    /// Prompt user to enter their API key
    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;
            }

            // Basic format validation
            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));
        }
    }

    /// Validate and test the API key
    async fn validate_and_test_key(&self, api_key: &SecureString) -> Result<()> {
        println!("\n{} {}", "".cyan().bold(), "Testing API key...".cyan());

        // Create a simple key provider for testing
        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);

        // Make a simple test request
        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
                )))
            }
        }
    }

    /// Print completion message
    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");
    }

    /// Quick setup without interactive prompts (for programmatic use)
    pub async fn quick_setup(
        api_key: String,
        provider: Option<Arc<dyn ApiKeyProvider>>,
    ) -> Result<()> {
        let secure_key = SecureString::new(api_key);

        // Validate the 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);

        // Test the key
        client
            .generate_content("Respond with just 'OK' if you receive this message.", None)
            .await?;

        // Store if provider available
        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");
    }
}