procyon 0.1.2

use std::collections::BTreeMap;
use std::path::PathBuf;

use color_eyre::{eyre::WrapErr, Result};
use serde::{Deserialize, Serialize};

use crate::config::AppConfig;

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CredentialEntry {
    api_key: String,
}

// Same 0600 rationale as `config.rs::restrict_to_owner`: this file holds live API keys, not just
// a reference to where they live.
#[cfg(unix)]
fn restrict_to_owner(path: &std::path::Path) -> Result<()> {
    use std::os::unix::fs::PermissionsExt;
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?;
    Ok(())
}

#[cfg(not(unix))]
fn restrict_to_owner(_path: &std::path::Path) -> Result<()> {
    Ok(())
}

/// API keys entered through `/login`, kept out of `config.toml` so the config file stays safe to
/// share or check in. Keyed by the same lowercase string `Provider`'s `Display` produces.
pub struct CredentialStore {
    path: PathBuf,
    entries: BTreeMap<String, CredentialEntry>,
}

impl CredentialStore {
    pub fn load(path: PathBuf) -> Result<Self> {
        let entries = if path.exists() {
            let content = std::fs::read_to_string(&path)
                .wrap_err_with(|| format!("Reading credentials at {}", path.display()))?;
            toml::from_str(&content)
                .wrap_err_with(|| format!("Invalid credentials at {}", path.display()))?
        } else {
            BTreeMap::new()
        };
        Ok(Self { path, entries })
    }

    /// The path real (non-test) callers use.
    pub fn load_default() -> Result<Self> {
        Self::load(AppConfig::config_dir()?.join("credentials.toml"))
    }

    pub fn get(&self, provider: &str) -> Option<&str> {
        self.entries.get(provider).map(|e| e.api_key.as_str())
    }

    pub fn set(&mut self, provider: &str, key: String) -> Result<()> {
        self.entries
            .insert(provider.to_string(), CredentialEntry { api_key: key });
        self.persist()
    }

    /// Returns whether an entry actually existed to remove.
    pub fn remove(&mut self, provider: &str) -> Result<bool> {
        let existed = self.entries.remove(provider).is_some();
        if existed {
            self.persist()?;
        }
        Ok(existed)
    }

    pub fn providers_with_keys(&self) -> impl Iterator<Item = &str> {
        self.entries.keys().map(|s| s.as_str())
    }

    fn persist(&self) -> Result<()> {
        if let Some(parent) = self.path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let content = toml::to_string_pretty(&self.entries)?;
        let tmp = self.path.with_extension("toml.tmp");
        std::fs::write(&tmp, &content)?;
        restrict_to_owner(&tmp)?;
        std::fs::rename(&tmp, &self.path)?;
        Ok(())
    }
}

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

    fn store_at(dir: &std::path::Path) -> CredentialStore {
        CredentialStore::load(dir.join("credentials.toml")).unwrap()
    }

    #[test]
    fn set_then_get_round_trips() {
        let temp = tempfile::tempdir().unwrap();
        let mut store = store_at(temp.path());

        store.set("anthropic", "sk-ant-123".to_string()).unwrap();

        assert_eq!(store.get("anthropic"), Some("sk-ant-123"));
    }

    #[test]
    fn a_missing_file_behaves_as_an_empty_store() {
        let temp = tempfile::tempdir().unwrap();
        let store = store_at(temp.path());

        assert_eq!(store.get("anthropic"), None);
        assert_eq!(store.providers_with_keys().count(), 0);
    }

    #[test]
    #[cfg(unix)]
    fn a_write_leaves_the_file_owner_only() {
        use std::os::unix::fs::PermissionsExt;

        let temp = tempfile::tempdir().unwrap();
        let mut store = store_at(temp.path());
        store.set("groq", "gsk-1".to_string()).unwrap();

        let path = temp.path().join("credentials.toml");
        let mode = std::fs::metadata(&path).unwrap().permissions().mode();
        assert_eq!(mode & 0o777, 0o600);
    }

    #[test]
    fn removing_an_unknown_provider_reports_no_entry() {
        let temp = tempfile::tempdir().unwrap();
        let mut store = store_at(temp.path());

        assert!(!store.remove("openai").unwrap());
    }

    #[test]
    fn removing_a_known_provider_removes_it_and_says_so() {
        let temp = tempfile::tempdir().unwrap();
        let mut store = store_at(temp.path());
        store.set("openai", "sk-1".to_string()).unwrap();

        assert!(store.remove("openai").unwrap());
        assert_eq!(store.get("openai"), None);
    }

    #[test]
    fn setting_an_existing_provider_again_overwrites_rather_than_duplicating() {
        let temp = tempfile::tempdir().unwrap();
        let mut store = store_at(temp.path());

        store.set("xai", "old-key".to_string()).unwrap();
        store.set("xai", "new-key".to_string()).unwrap();

        assert_eq!(store.get("xai"), Some("new-key"));
        assert_eq!(store.providers_with_keys().count(), 1);
    }

    #[test]
    fn a_stored_key_survives_a_reload_from_disk() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("credentials.toml");
        let mut store = CredentialStore::load(path.clone()).unwrap();
        store.set("deepseek", "dsk-1".to_string()).unwrap();

        let reloaded = CredentialStore::load(path).unwrap();
        assert_eq!(reloaded.get("deepseek"), Some("dsk-1"));
    }
}