gyandoot-v2/crates/gyandoot-tui/src/setup.rs

//! Interactive setup wizard — configures providers and initializes GyanDoot.
//!
//! Flow:
//! 1. Print compact banner header
//! 2. Provider selection (numbered menu)
//! 3. Per-provider config (Base URL, API Key, Model, Temperature, Max Tokens)
//! 4. Summary display before saving
//! 5. Connection test (/models endpoint)
//! 6. Write config.toml to ~/.config/gyandoot/config.toml

use std::fmt;
use std::path::PathBuf;

use serde::{Deserialize, Serialize};

use crate::banner;
use crate::terminal;

/// Supported LLM providers.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProviderKind {
    OpenAI,
    Anthropic,
    Ollama,
    GoogleGemini,
    DeepSeek,
    OpenRouter,
    Groq,
}

impl ProviderKind {
    /// All known providers in display order.
    pub fn all() -> &'static [ProviderKind] {
        &[
            ProviderKind::OpenAI,
            ProviderKind::Anthropic,
            ProviderKind::Ollama,
            ProviderKind::GoogleGemini,
            ProviderKind::DeepSeek,
            ProviderKind::OpenRouter,
            ProviderKind::Groq,
        ]
    }

    /// Display name for menus.
    pub fn display_name(&self) -> &'static str {
        match self {
            ProviderKind::OpenAI => "OpenAI",
            ProviderKind::Anthropic => "Anthropic",
            ProviderKind::Ollama => "Ollama",
            ProviderKind::GoogleGemini => "Google Gemini",
            ProviderKind::DeepSeek => "DeepSeek",
            ProviderKind::OpenRouter => "OpenRouter",
            ProviderKind::Groq => "Groq",
        }
    }

    /// Default base URL for the provider.
    pub fn default_base_url(&self) -> &'static str {
        match self {
            ProviderKind::OpenAI => "https://api.openai.com/v1",
            ProviderKind::Anthropic => "https://api.anthropic.com/v1",
            ProviderKind::Ollama => "http://localhost:11434",
            ProviderKind::GoogleGemini => "https://generativelanguage.googleapis.com/v1",
            ProviderKind::DeepSeek => "https://api.deepseek.com/v1",
            ProviderKind::OpenRouter => "https://openrouter.ai/api/v1",
            ProviderKind::Groq => "https://api.groq.com/openai/v1",
        }
    }

    /// Default model name.
    pub fn default_model(&self) -> &'static str {
        match self {
            ProviderKind::OpenAI => "gpt-4o",
            ProviderKind::Anthropic => "claude-sonnet-4-20250514",
            ProviderKind::Ollama => "llama3",
            ProviderKind::GoogleGemini => "gemini-2.0-flash",
            ProviderKind::DeepSeek => "deepseek-chat",
            ProviderKind::OpenRouter => "openai/gpt-4o",
            ProviderKind::Groq => "llama3-70b-8192",
        }
    }

    /// Default temperature.
    pub fn default_temperature(&self) -> f64 {
        0.7
    }

    /// Default max tokens.
    pub fn default_max_tokens(&self) -> u32 {
        4096
    }

    /// Provider-specific API key env var hint.
    pub fn api_key_env_hint(&self) -> &'static str {
        match self {
            ProviderKind::OpenAI => "OPENAI_API_KEY",
            ProviderKind::Anthropic => "ANTHROPIC_API_KEY",
            ProviderKind::Ollama => "OLLAMA_API_KEY (optional)",
            ProviderKind::GoogleGemini => "GEMINI_API_KEY",
            ProviderKind::DeepSeek => "DEEPSEEK_API_KEY",
            ProviderKind::OpenRouter => "OPENROUTER_API_KEY",
            ProviderKind::Groq => "GROQ_API_KEY",
        }
    }
}

impl fmt::Display for ProviderKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.display_name())
    }
}

/// Per-provider configuration collected during setup.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderConfig {
    /// Which provider this config is for.
    pub kind: ProviderKind,
    /// API base URL.
    pub base_url: String,
    /// API key (may be empty for local Ollama).
    pub api_key: String,
    /// Model identifier.
    pub model: String,
    /// Temperature parameter (0.0–2.0).
    pub temperature: f64,
    /// Max tokens per response.
    pub max_tokens: u32,
    /// Ollama-specific: merge system messages.
    #[serde(default)]
    pub merge_system: bool,
    /// Ollama-specific: disable tool calling.
    #[serde(default)]
    pub no_tools: bool,
}

impl ProviderConfig {
    /// Create a new provider config with defaults.
    pub fn new(kind: ProviderKind) -> Self {
        Self {
            kind,
            base_url: kind.default_base_url().to_string(),
            api_key: String::new(),
            model: kind.default_model().to_string(),
            temperature: kind.default_temperature(),
            max_tokens: kind.default_max_tokens(),
            merge_system: false,
            no_tools: false,
        }
    }
}

/// The result of the setup wizard.
#[derive(Debug, Serialize, Deserialize)]
pub struct WizardResult {
    /// Selected provider kind.
    pub provider: ProviderKind,
    /// Full provider config.
    pub config: ProviderConfig,
    /// Whether the connection test passed.
    pub connection_ok: bool,
    /// Path where config was written.
    pub config_path: PathBuf,
}

/// Default config directory.
fn config_dir() -> PathBuf {
    dirs::home_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join(".config")
        .join("gyandoot")
}

/// Run the full interactive setup wizard.
///
/// Returns the wizard result on success, or an error if the user cancels or
/// something goes wrong during write.
pub fn run_setup_wizard() -> anyhow::Result<Option<WizardResult>> {
    // Step 1: Print compact banner
    println!();
    println!("{}", banner::render_fallback_colored());
    println!();
    terminal::section_header("Setup Wizard");

    // Step 2: Provider selection
    println!("Select your LLM provider:");
    let provider_names: Vec<&str> = ProviderKind::all()
        .iter()
        .map(|p| p.display_name())
        .collect();

    let choice = terminal::selection_menu("Choose a provider:", &provider_names);
    let choice = match choice {
        Some(idx) => ProviderKind::all()[idx],
        None => {
            terminal::show_status(terminal::StatusKind::Info, "Setup cancelled.");
            return Ok(None);
        }
    };

    println!();
    terminal::show_status(
        terminal::StatusKind::Info,
        &format!("Configuring provider: {}", choice.display_name()),
    );

    // Step 3: Per-provider config
    let mut config = ProviderConfig::new(choice);

    // Base URL
    println!();
    let url = terminal::text_input("Base URL", Some(&config.base_url));
    match url {
        Some(u) if !u.trim().is_empty() => config.base_url = u.trim().to_string(),
        Some(_) => {} // keep default
        None => return Ok(None),
    }

    // API Key (with env hint)
    let api_key_hint = choice.api_key_env_hint();
    println!("  (Hint: set via {} env var or enter now)", api_key_hint);
    let key = terminal::text_input("API Key", None);
    match key {
        Some(k) => config.api_key = k.trim().to_string(),
        None => return Ok(None),
    }

    // Model
    let model = terminal::text_input("Model", Some(&config.model));
    match model {
        Some(m) if !m.trim().is_empty() => config.model = m.trim().to_string(),
        Some(_) => {} // keep default
        None => return Ok(None),
    }

    // Temperature
    let temp_input = terminal::text_input("Temperature (0.0-2.0)", Some("0.7"));
    match temp_input {
        Some(t) if !t.trim().is_empty() => {
            if let Ok(v) = t.trim().parse::<f64>() {
                config.temperature = v.clamp(0.0, 2.0);
            }
        }
        Some(_) => {} // keep default
        None => return Ok(None),
    }

    // Max Tokens
    let tokens_input = terminal::text_input("Max Tokens", Some("4096"));
    match tokens_input {
        Some(t) if !t.trim().is_empty() => {
            if let Ok(v) = t.trim().parse::<u32>() {
                config.max_tokens = v;
            }
        }
        Some(_) => {} // keep default
        None => return Ok(None),
    }

    // Auto-detect Ollama localhost settings
    if choice == ProviderKind::Ollama && config.base_url.contains("localhost") {
        config.merge_system = true;
        config.no_tools = true;
        println!();
        terminal::show_status(
            terminal::StatusKind::Info,
            "Local Ollama detected: merge_system=true, no_tools=true",
        );
    }

    // Step 4: Summary
    println!();
    terminal::section_header("Configuration Summary");
    println!("  Provider:    {}", config.kind.display_name());
    println!("  Base URL:    {}", config.base_url);
    println!(
        "  API Key:     {}",
        if config.api_key.is_empty() {
            "(none)".to_string()
        } else {
            format!("{}…{}", &config.api_key[..3.min(config.api_key.len())], &config.api_key[config.api_key.len().saturating_sub(4)..])
        }
    );
    println!("  Model:       {}", config.model);
    println!("  Temperature: {}", config.temperature);
    println!("  Max Tokens:  {}", config.max_tokens);
    if config.merge_system {
        println!("  Merge System: true");
    }
    if config.no_tools {
        println!("  No Tools:    true");
    }

    println!();
    if !terminal::confirm_prompt("Save this configuration?", Some(true)) {
        terminal::show_status(terminal::StatusKind::Info, "Setup cancelled.");
        return Ok(None);
    }

    // Step 5: Connection test
    println!();
    terminal::show_status(
        terminal::StatusKind::Info,
        "Running connection test...",
    );

    let test_passed = test_provider_connection(&config);

    if test_passed {
        terminal::show_status(
            terminal::StatusKind::Success,
            "Connection successful!",
        );
    } else {
        terminal::show_status(
            terminal::StatusKind::Warning,
            "Connection test failed — config will still be saved.",
        );
        println!("  (You can retry later or adjust settings manually)");
    }

    // Step 6: Save config
    println!();
    if !terminal::confirm_prompt("Write config file?", Some(true)) {
        terminal::show_status(terminal::StatusKind::Info, "Setup cancelled.");
        return Ok(None);
    }

    let config_path = write_config_file(&config)?;

    terminal::show_status(
        terminal::StatusKind::Success,
        &format!("Configuration saved to {}", config_path.display()),
    );

    Ok(Some(WizardResult {
        provider: choice,
        config,
        connection_ok: test_passed,
        config_path,
    }))
}

/// Test provider connectivity by calling the /models endpoint.
fn test_provider_connection(config: &ProviderConfig) -> bool {
    if config.api_key.is_empty() && config.kind != ProviderKind::Ollama {
        return false;
    }

    let client = reqwest::blocking::Client::builder()
        .timeout(std::time::Duration::from_secs(10))
        .build()
        .ok();

    let client = match client {
        Some(c) => c,
        None => return false,
    };

    let url = format!("{}/models", config.base_url.trim_end_matches('/'));

    let request = client.get(&url);
    let request = if !config.api_key.is_empty() {
        request.header("Authorization", format!("Bearer {}", config.api_key))
    } else {
        request
    };

    match request.send() {
        Ok(resp) => resp.status().is_success() || resp.status().as_u16() == 401 || resp.status().as_u16() == 403,
        Err(_) => false,
    }
}

/// Write the provider configuration to ~/.config/gyandoot/config.toml.
fn write_config_file(config: &ProviderConfig) -> anyhow::Result<PathBuf> {
    let dir = config_dir();
    std::fs::create_dir_all(&dir)?;

    let path = dir.join("config.toml");

    let toml_content = toml::to_string_pretty(&SerializedConfig::from(config))?;
    std::fs::write(&path, &toml_content)?;

    Ok(path)
}

/// Serializable wrapper for TOML output.
#[derive(Serialize)]
struct SerializedConfig {
    provider: SerializedProvider,
    agent: SerializedAgent,
}

#[derive(Serialize)]
struct SerializedProvider {
    kind: String,
    base_url: String,
    api_key: String,
    model: String,
    temperature: f64,
    max_tokens: u32,
    #[serde(skip_serializing_if = "is_false")]
    merge_system: bool,
    #[serde(skip_serializing_if = "is_false")]
    no_tools: bool,
}

#[derive(Serialize)]
struct SerializedAgent {
    name: String,
}

fn is_false(v: &bool) -> bool {
    !v
}

impl From<&ProviderConfig> for SerializedConfig {
    fn from(c: &ProviderConfig) -> Self {
        let kind_str = match c.kind {
            ProviderKind::OpenAI => "openai",
            ProviderKind::Anthropic => "anthropic",
            ProviderKind::Ollama => "ollama",
            ProviderKind::GoogleGemini => "google-gemini",
            ProviderKind::DeepSeek => "deepseek",
            ProviderKind::OpenRouter => "openrouter",
            ProviderKind::Groq => "groq",
        };

        SerializedConfig {
            provider: SerializedProvider {
                kind: kind_str.to_string(),
                base_url: c.base_url.clone(),
                api_key: c.api_key.clone(),
                model: c.model.clone(),
                temperature: c.temperature,
                max_tokens: c.max_tokens,
                merge_system: c.merge_system,
                no_tools: c.no_tools,
            },
            agent: SerializedAgent {
                name: "gyandoot".to_string(),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_provider_kind_all() {
        let all = ProviderKind::all();
        assert!(!all.is_empty());
        assert_eq!(all.len(), 7);
    }

    #[test]
    fn test_provider_display_names() {
        assert_eq!(ProviderKind::OpenAI.display_name(), "OpenAI");
        assert_eq!(ProviderKind::Ollama.display_name(), "Ollama");
    }

    #[test]
    fn test_provider_defaults() {
        for kind in ProviderKind::all() {
            assert!(!kind.default_base_url().is_empty());
            assert!(!kind.default_model().is_empty());
            assert!(kind.default_temperature() > 0.0);
            assert!(kind.default_max_tokens() > 0);
        }
    }

    #[test]
    fn test_provider_config_new() {
        let cfg = ProviderConfig::new(ProviderKind::OpenAI);
        assert_eq!(cfg.kind, ProviderKind::OpenAI);
        assert!(cfg.base_url.contains("openai"));
        assert_eq!(cfg.model, "gpt-4o");
        assert_eq!(cfg.temperature, 0.7);
        assert_eq!(cfg.max_tokens, 4096);
    }

    #[test]
    fn test_config_dir_has_gyandoot() {
        let dir = config_dir();
        assert!(dir.to_string_lossy().contains("gyandoot"));
    }

    #[test]
    fn test_serialized_config_roundtrip() {
        let cfg = ProviderConfig::new(ProviderKind::Anthropic);
        let serialized = SerializedConfig::from(&cfg);
        let toml_str = toml::to_string_pretty(&serialized).unwrap();
        assert!(toml_str.contains("anthropic"));
        assert!(toml_str.contains("base_url"));
        assert!(toml_str.contains("api_key"));
    }

    #[test]
    fn test_ollama_autodetect() {
        let mut cfg = ProviderConfig::new(ProviderKind::Ollama);
        cfg.base_url = "http://localhost:11434".to_string();
        cfg.merge_system = true;
        cfg.no_tools = true;
        assert!(cfg.merge_system);
        assert!(cfg.no_tools);
    }
}