QuiCommit/src/utils/keyring.rs

use anyhow::{Context, Result, bail};
use std::env;

const SERVICE_NAME: &str = "quicommit";
const ENV_API_KEY: &str = "QUICOMMIT_API_KEY";

const PAT_SERVICE_PREFIX: &str = "quicommit/pat";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyringStatus {
    Available,
    Unavailable,
}

pub struct KeyringManager {
    status: KeyringStatus,
}

impl KeyringManager {
    pub fn new() -> Self {
        let status = Self::check_keyring_availability();
        Self { status }
    }

    pub fn check_keyring_availability() -> KeyringStatus {
        #[cfg(target_os = "windows")]
        {
            KeyringStatus::Available
        }
        #[cfg(target_os = "macos")]
        {
            KeyringStatus::Available
        }
        #[cfg(target_os = "linux")]
        {
            Self::check_linux_keyring()
        }
        #[cfg(not(any(target_os = "windows", target_os = "macos", target_os = "linux")))]
        {
            KeyringStatus::Unavailable
        }
    }

    #[cfg(target_os = "linux")]
    fn check_linux_keyring() -> KeyringStatus {
        use std::path::Path;

        let has_dbus = Path::new("/usr/bin/dbus-daemon").exists()
            || Path::new("/bin/dbus-daemon").exists()
            || env::var("DBUS_SESSION_BUS_ADDRESS").is_ok();

        let has_keyring = Path::new("/usr/bin/gnome-keyring-daemon").exists()
            || Path::new("/usr/bin/gnome-keyring").exists()
            || Path::new("/usr/bin/kwalletd5").exists()
            || Path::new("/usr/bin/kwalletd6").exists()
            || env::var("SECRET_SERVICE").is_ok();

        if has_dbus && has_keyring {
            KeyringStatus::Available
        } else {
            KeyringStatus::Unavailable
        }
    }

    pub fn status(&self) -> KeyringStatus {
        self.status
    }

    pub fn is_available(&self) -> bool {
        self.status == KeyringStatus::Available
    }

    pub fn store_api_key(&self, provider: &str, api_key: &str) -> Result<()> {
        if !self.is_available() {
            bail!("Keyring is not available on this system");
        }

        let entry = keyring::Entry::new(SERVICE_NAME, provider)
            .context("Failed to create keyring entry")?;

        entry
            .set_password(api_key)
            .context("Failed to store API key")?;

        Ok(())
    }

    pub fn get_api_key(&self, provider: &str) -> Result<Option<String>> {
        if let Ok(key) = env::var(ENV_API_KEY)
            && !key.is_empty()
        {
            return Ok(Some(key));
        }

        if !self.is_available() {
            return Ok(None);
        }

        let entry = keyring::Entry::new(SERVICE_NAME, provider)
            .context("Failed to create keyring entry")?;

        match entry.get_password() {
            Ok(key) => Ok(Some(key)),
            Err(keyring::Error::NoEntry) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    pub fn delete_api_key(&self, provider: &str) -> Result<()> {
        if !self.is_available() {
            bail!("Keyring is not available on this system");
        }

        let entry = keyring::Entry::new(SERVICE_NAME, provider)
            .context("Failed to create keyring entry")?;

        entry
            .delete_credential()
            .context("Failed to delete API key")?;

        Ok(())
    }

    pub fn has_api_key(&self, provider: &str) -> bool {
        self.get_api_key(provider).unwrap_or(None).is_some()
    }

    fn make_pat_service_name(profile_name: &str) -> String {
        format!("{}/{}", PAT_SERVICE_PREFIX, profile_name)
    }

    pub fn store_pat(
        &self,
        profile_name: &str,
        user_email: &str,
        service: &str,
        token: &str,
    ) -> Result<()> {
        if !self.is_available() {
            bail!("Keyring is not available on this system");
        }

        let keyring_service = Self::make_pat_service_name(profile_name);
        let keyring_user = format!("{}:{}", user_email, service);

        let entry = keyring::Entry::new(&keyring_service, &keyring_user)
            .context("Failed to create keyring entry for PAT")?;

        entry
            .set_password(token)
            .context("Failed to store PAT in keyring")?;

        Ok(())
    }

    pub fn get_pat(
        &self,
        profile_name: &str,
        user_email: &str,
        service: &str,
    ) -> Result<Option<String>> {
        if !self.is_available() {
            return Ok(None);
        }

        let keyring_service = Self::make_pat_service_name(profile_name);
        let keyring_user = format!("{}:{}", user_email, service);

        let entry = keyring::Entry::new(&keyring_service, &keyring_user)
            .context("Failed to create keyring entry for PAT")?;

        match entry.get_password() {
            Ok(token) => Ok(Some(token)),
            Err(keyring::Error::NoEntry) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    pub fn delete_pat(&self, profile_name: &str, user_email: &str, service: &str) -> Result<()> {
        if !self.is_available() {
            bail!("Keyring is not available on this system");
        }

        let keyring_service = Self::make_pat_service_name(profile_name);
        let keyring_user = format!("{}:{}", user_email, service);

        let entry = keyring::Entry::new(&keyring_service, &keyring_user)
            .context("Failed to create keyring entry for PAT")?;

        entry
            .delete_credential()
            .context("Failed to delete PAT from keyring")?;

        Ok(())
    }

    pub fn has_pat(&self, profile_name: &str, user_email: &str, service: &str) -> bool {
        self.get_pat(profile_name, user_email, service)
            .unwrap_or(None)
            .is_some()
    }

    pub fn delete_all_pats_for_profile(
        &self,
        profile_name: &str,
        user_email: &str,
        services: &[String],
    ) -> Result<()> {
        for service in services {
            let _ = self.delete_pat(profile_name, user_email, service);
        }
        Ok(())
    }

    pub fn get_status_message(&self) -> String {
        match self.status {
            KeyringStatus::Available => {
                #[cfg(target_os = "windows")]
                {
                    "Windows Credential Manager is available".to_string()
                }
                #[cfg(target_os = "macos")]
                {
                    "macOS Keychain is available".to_string()
                }
                #[cfg(target_os = "linux")]
                {
                    "Linux secret service is available".to_string()
                }
                #[cfg(not(any(target_os = "windows", target_os = "macos", target_os = "linux")))]
                {
                    "Keyring is available".to_string()
                }
            }
            KeyringStatus::Unavailable => {
                "Keyring is not available. Set QUICOMMIT_API_KEY environment variable.".to_string()
            }
        }
    }
}

impl Default for KeyringManager {
    fn default() -> Self {
        Self::new()
    }
}

pub fn get_default_base_url(provider: &str) -> &'static str {
    match provider {
        "openai" => "https://api.openai.com/v1",
        "anthropic" => "https://api.anthropic.com/v1",
        "kimi" => "https://api.moonshot.cn/v1",
        "deepseek" => "https://api.deepseek.com/v1",
        "openrouter" => "https://openrouter.ai/api/v1",
        "ollama" => "http://localhost:11434",
        _ => "",
    }
}

pub fn get_default_model(provider: &str) -> &'static str {
    match provider {
        "openai" => "gpt-4",
        "anthropic" => "claude-3-sonnet-20240229",
        "kimi" => "kimi-k2.6",
        "deepseek" => "deepseek-v4-flash",
        "openrouter" => "openai/gpt-3.5-turbo",
        "ollama" => "llama2",
        _ => "",
    }
}

pub fn get_supported_providers() -> &'static [&'static str] {
    &[
        "ollama",
        "openai",
        "anthropic",
        "kimi",
        "deepseek",
        "openrouter",
    ]
}

pub fn provider_needs_api_key(provider: &str) -> bool {
    provider != "ollama"
}

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

    #[test]
    fn test_get_default_base_url() {
        assert_eq!(get_default_base_url("openai"), "https://api.openai.com/v1");
        assert_eq!(
            get_default_base_url("anthropic"),
            "https://api.anthropic.com/v1"
        );
        assert_eq!(get_default_base_url("kimi"), "https://api.moonshot.cn/v1");
        assert_eq!(
            get_default_base_url("deepseek"),
            "https://api.deepseek.com/v1"
        );
        assert_eq!(
            get_default_base_url("openrouter"),
            "https://openrouter.ai/api/v1"
        );
        assert_eq!(get_default_base_url("ollama"), "http://localhost:11434");
    }

    #[test]
    fn test_get_default_model() {
        assert_eq!(get_default_model("openai"), "gpt-4");
        assert_eq!(get_default_model("anthropic"), "claude-3-sonnet-20240229");
        assert_eq!(get_default_model("ollama"), "llama2");
    }

    #[test]
    fn test_provider_needs_api_key() {
        assert!(provider_needs_api_key("openai"));
        assert!(provider_needs_api_key("anthropic"));
        assert!(!provider_needs_api_key("ollama"));
    }
}