procyon 0.1.2

use color_eyre::{eyre::bail, eyre::WrapErr, Result};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

use crate::credentials::CredentialStore;

// Must stay lowercase to match `Display`/`FromStr`, which is the spelling the config file and
// the `/theme` command use.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Theme {
    Light,
    Dark,
}

impl std::fmt::Display for Theme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Theme::Light => write!(f, "light"),
            Theme::Dark => write!(f, "dark"),
        }
    }
}

impl std::str::FromStr for Theme {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "light" => Ok(Theme::Light),
            "dark" => Ok(Theme::Dark),
            _ => Err(format!("Unknown theme: {}", s)),
        }
    }
}

/// Which provider to talk to.
///
/// The variants that matter to the code are the two *protocols* — Anthropic's `/v1/messages` and
/// the OpenAI `/chat/completions` dialect. Everything else is a named profile: a preset endpoint
/// and credential variable over the OpenAI protocol, so `provider = "groq"` is all the config
/// needs rather than a base URL the user has to look up. Adding a vendor is one table row.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "kebab-case")]
pub enum Provider {
    Anthropic,
    OpenAi,
    Deepseek,
    Groq,
    Openrouter,
    Cerebras,
    Fireworks,
    Togetherai,
    Xai,
    Ollama,
    LmStudio,
    /// Any other endpoint speaking the OpenAI dialect; requires an explicit `base_url`.
    OpenAiCompatible,
}

/// The preset behind a named provider: where to post and which variable holds the credential.
struct Profile {
    base_url: &'static str,
    key_env: &'static str,
}

impl Provider {
    /// `None` for Anthropic, whose endpoint lives in its own client, and for the generic
    /// `openai-compatible` escape hatch, which has nothing to preset.
    fn profile(&self) -> Option<Profile> {
        let (base_url, key_env) = match self {
            Provider::Anthropic | Provider::OpenAiCompatible => return None,
            Provider::OpenAi => ("https://api.openai.com/v1", "OPENAI_API_KEY"),
            Provider::Deepseek => ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY"),
            Provider::Groq => ("https://api.groq.com/openai/v1", "GROQ_API_KEY"),
            Provider::Openrouter => ("https://openrouter.ai/api/v1", "OPENROUTER_API_KEY"),
            Provider::Cerebras => ("https://api.cerebras.ai/v1", "CEREBRAS_API_KEY"),
            Provider::Fireworks => ("https://api.fireworks.ai/inference/v1", "FIREWORKS_API_KEY"),
            Provider::Togetherai => ("https://api.together.xyz/v1", "TOGETHER_API_KEY"),
            Provider::Xai => ("https://api.x.ai/v1", "XAI_API_KEY"),
            // Served from this machine: the endpoint is a default rather than a vendor address,
            // and the credential variable is nominal since neither checks it.
            Provider::Ollama => ("http://localhost:11434/v1", "OLLAMA_API_KEY"),
            Provider::LmStudio => ("http://localhost:1234/v1", "LM_STUDIO_API_KEY"),
        };
        Some(Profile { base_url, key_env })
    }

    /// Whether this provider speaks Anthropic's `/v1/messages` rather than the OpenAI dialect.
    pub fn is_anthropic(&self) -> bool {
        matches!(self, Provider::Anthropic)
    }

    /// The variable consulted when the config names no `api_key_env` of its own.
    pub fn default_key_env(&self) -> &'static str {
        match self {
            Provider::Anthropic => "ANTHROPIC_API_KEY",
            other => other
                .profile()
                .map(|p| p.key_env)
                .unwrap_or("OPENAI_API_KEY"),
        }
    }

    /// Served from this machine, so nothing checks a credential for it.
    pub fn is_local(&self) -> bool {
        matches!(self, Provider::Ollama | Provider::LmStudio)
    }

    /// Every named provider, for pickers that need to enumerate them (the wizard's provider step,
    /// `/providers`). Excludes `OpenAiCompatible`, which needs a `base_url` a list cannot supply.
    pub const ALL: [Provider; 11] = [
        Provider::Anthropic,
        Provider::OpenAi,
        Provider::Deepseek,
        Provider::Groq,
        Provider::Openrouter,
        Provider::Cerebras,
        Provider::Fireworks,
        Provider::Togetherai,
        Provider::Xai,
        Provider::Ollama,
        Provider::LmStudio,
    ];
}

impl std::fmt::Display for Provider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Provider::Anthropic => write!(f, "anthropic"),
            Provider::OpenAi => write!(f, "openai"),
            Provider::Deepseek => write!(f, "deepseek"),
            Provider::Groq => write!(f, "groq"),
            Provider::Openrouter => write!(f, "openrouter"),
            Provider::Cerebras => write!(f, "cerebras"),
            Provider::Fireworks => write!(f, "fireworks"),
            Provider::Togetherai => write!(f, "togetherai"),
            Provider::Xai => write!(f, "xai"),
            Provider::Ollama => write!(f, "ollama"),
            Provider::LmStudio => write!(f, "lmstudio"),
            Provider::OpenAiCompatible => write!(f, "openai-compatible"),
        }
    }
}

impl std::str::FromStr for Provider {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "anthropic" => Ok(Provider::Anthropic),
            "openai" => Ok(Provider::OpenAi),
            "deepseek" => Ok(Provider::Deepseek),
            "groq" => Ok(Provider::Groq),
            "openrouter" => Ok(Provider::Openrouter),
            "cerebras" => Ok(Provider::Cerebras),
            "fireworks" => Ok(Provider::Fireworks),
            "togetherai" => Ok(Provider::Togetherai),
            "xai" => Ok(Provider::Xai),
            "ollama" => Ok(Provider::Ollama),
            "lmstudio" | "lm-studio" | "lm_studio" => Ok(Provider::LmStudio),
            "openai-compatible" | "openai_compatible" | "openaicompatible" => {
                Ok(Provider::OpenAiCompatible)
            }
            _ => Err(format!(
                "Unknown provider: {}. Use: anthropic, openai, deepseek, groq, openrouter, \
                 cerebras, fireworks, togetherai, xai, ollama, lmstudio, openai-compatible",
                s
            )),
        }
    }
}

impl Provider {
    /// Popular models available through this provider.
    pub fn suggested_models(&self) -> Vec<&'static str> {
        match self {
            Provider::Anthropic => vec!["claude-sonnet-5", "claude-opus-5", "claude-haiku"],
            Provider::OpenAi => vec!["gpt-4.1", "gpt-4o", "o3", "o4-mini"],
            Provider::Deepseek => vec!["deepseek-chat", "deepseek-reasoner"],
            Provider::Groq => vec![
                "llama-3.3-70b-versatile",
                "llama-3.1-8b-instant",
                "mixtral-8x7b-32768",
            ],
            Provider::Openrouter => vec![
                "anthropic/claude-sonnet-5",
                "openai/gpt-4.1",
                "google/gemini-2.5-pro",
                "meta-llama/llama-3.3-70b-instruct",
            ],
            Provider::Cerebras => vec!["llama-3.3-70b", "llama-3.1-8b"],
            Provider::Fireworks => vec![
                "accounts/fireworks/models/llama-v3p3-70b-instruct",
                "accounts/fireworks/models/mixtral-8x22b-instruct",
            ],
            Provider::Togetherai => vec![
                "meta-llama/Llama-3.3-70B-Instruct-Turbo",
                "deepseek-ai/DeepSeek-V3",
                "Qwen/Qwen2.5-72B-Instruct-Turbo",
            ],
            Provider::Xai => vec!["grok-3", "grok-3-mini"],
            // Ordered by tool-calling ability, not by size or recency, because the first entry is
            // also `default_model()` — what a user who just picks "ollama" ends up running. The
            // 3B and 4B tags used to lead, and a model that small does not reliably emit a tool
            // call at all: it answers a request to run the tests by stating that no test runner
            // exists, which reads as a missing feature rather than as a model that cannot reach
            // one. 8B is the smallest size that behaves. `codellama` is gone for the same reason —
            // it predates tool calling entirely.
            Provider::Ollama => vec![
                "llama3.1:8b",
                "qwen2.5-coder:7b",
                "qwen3:8b",
                "mistral",
                "deepseek-coder-v2",
            ],
            Provider::LmStudio => vec![
                "lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF",
                "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
            ],
            Provider::OpenAiCompatible => vec![],
        }
    }

    /// The model to select when this provider is chosen without naming one.
    ///
    /// Switching provider used to keep whatever model was already active, producing pairs no
    /// endpoint serves — `ollama` still asking for `claude-sonnet-5`. Empty only for the generic
    /// OpenAI-compatible provider, where nothing can be guessed and the operator has to say which
    /// model the endpoint serves.
    pub fn default_model(&self) -> &'static str {
        self.suggested_models().first().copied().unwrap_or("")
    }
}

/// One MCP server, as declared in `config.toml`. Either remote over HTTP (`url`) or a local
/// process speaking stdio (`command`), never both.
///
/// Declared here rather than in `mcp` because it is part of the config file's shape: config
/// described its own schema in terms of a protocol module, so `AppConfig` could not be read
/// without pulling in an MCP client. `mcp` interprets these fields — see `transport`.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct McpServerConfig {
    /// Prefixes the server's tool names, keeping generic ones like `search` unambiguous.
    pub name: String,
    /// Streamable HTTP endpoint.
    pub url: Option<String>,
    /// Name of the environment variable holding the bearer credential. The token itself is never
    /// stored in the config file.
    pub token_env: Option<String>,
    /// Executable to launch for a stdio server. It inherits this process's environment.
    pub command: Option<String>,
    pub args: Vec<String>,
    /// Set to "oauth" for a server that requires browser sign-in. Mutually exclusive with
    /// `token_env`.
    pub auth: Option<String>,
}

// Every field defaults, so a config may set only what it wants to override instead of having to
// restate the whole struct.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct AppConfig {
    pub anthropic_api_key: Option<String>,
    pub default_network: String,
    pub theme: Theme,
    pub provider: Provider,
    pub default_model: String,
    /// Required by `openai-compatible`; ignored by `anthropic`, whose endpoint is fixed.
    pub base_url: Option<String>,
    /// Name of the environment variable holding the credential, so the key itself never has to
    /// be written into the config file.
    pub api_key_env: Option<String>,
    pub max_tokens: u32,
    /// Whether tools may sign and submit on the public network.
    ///
    /// Off by default, and deliberately a setting rather than a prompt: a mainnet operation spends
    /// real funds, and the agent must not be able to grant itself permission mid-conversation. The
    /// environment variable `PROCYON_ALLOW_MAINNET` overrides this for a single run.
    pub allow_mainnet: bool,
    pub mcp_servers: Vec<McpServerConfig>,
}

impl Default for AppConfig {
    fn default() -> Self {
        Self {
            anthropic_api_key: None,
            default_network: "testnet".to_string(),
            theme: Theme::Dark,
            provider: Provider::Anthropic,
            default_model: "claude-sonnet-5".to_string(),
            base_url: None,
            api_key_env: None,
            max_tokens: 4096,
            allow_mainnet: false,
            mcp_servers: Vec::new(),
        }
    }
}

// The config holds the API key, so it must not be readable by other local users.
#[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(())
}

impl AppConfig {
    pub fn config_dir() -> Result<PathBuf> {
        let config_dir = dirs::config_dir()
            .ok_or_else(|| color_eyre::eyre::eyre!("Failed to get config directory"))?
            .join("procyon");
        Ok(config_dir)
    }

    pub fn config_path() -> Result<PathBuf> {
        Ok(Self::config_dir()?.join("config.toml"))
    }

    pub fn load() -> Result<Self> {
        let path = Self::config_path()?;
        if !path.exists() {
            return Ok(Self::default());
        }
        let content = std::fs::read_to_string(&path)?;
        let config: Self = toml::from_str(&content)
            .wrap_err_with(|| format!("Invalid config at {}", path.display()))?;
        Ok(config)
    }

    pub fn save(&self) -> Result<()> {
        let config_dir = Self::config_dir()?;
        std::fs::create_dir_all(&config_dir)?;
        let path = Self::config_path()?;
        let content = toml::to_string_pretty(self)?;
        let tmp = path.with_extension("toml.tmp");
        std::fs::write(&tmp, &content)?;
        restrict_to_owner(&tmp)?;
        std::fs::rename(&tmp, &path)?;
        Ok(())
    }

    pub fn is_first_run() -> bool {
        match Self::config_path() {
            Ok(path) => !path.exists(),
            Err(_) => false,
        }
    }

    pub fn key_env_var(&self) -> &str {
        self.api_key_env
            .as_deref()
            .unwrap_or_else(|| self.provider.default_key_env())
    }

    /// A model served from this machine needs no credential, and demanding one would block the
    /// most convenient way to run without any provider account at all.
    fn endpoint_is_local(&self) -> bool {
        match self.resolve_base_url() {
            Ok(url) => {
                url.contains("://localhost")
                    || url.contains("://127.0.0.1")
                    || url.contains("://[::1]")
            }
            Err(_) => false,
        }
    }

    pub fn get_api_key(&self) -> Result<String> {
        // The real store, resolved fresh each call: `/login` and `/logout` may have edited the
        // file since the last time a key was needed. A store that fails to load (e.g. malformed
        // TOML) is treated as absent rather than aborting the whole lookup — the env var and local
        // bypass below still deserve a chance.
        self.get_api_key_with_store(CredentialStore::load_default().ok().as_ref())
    }

    /// Split out so tests can hand in a `CredentialStore` pointed at a tempdir instead of the real
    /// `~/.config/procyon/credentials.toml`.
    fn get_api_key_with_store(&self, store: Option<&CredentialStore>) -> Result<String> {
        // Retained for configs written before `provider` existed, where the inline key could only
        // ever have been an Anthropic one.
        if self.provider.is_anthropic() {
            if let Some(key) = &self.anthropic_api_key {
                return Ok(key.clone());
            }
        }

        if let Some(key) = store
            .and_then(|s| s.get(&self.provider.to_string()))
            .filter(|k| !k.is_empty())
        {
            return Ok(key.to_string());
        }

        let var = self.key_env_var();
        if let Some(key) = std::env::var(var).ok().filter(|k| !k.is_empty()) {
            return Ok(key);
        }

        if self.endpoint_is_local() {
            return Ok(String::new());
        }

        bail!(
            "No API key found for provider {:?}. Set {} in the environment, run `/login {} <key>` \
             inside procyon, or point `api_key_env` at the variable that holds it in \
             ~/.config/procyon/config.toml",
            self.provider,
            var,
            self.provider
        )
    }

    /// The endpoint the OpenAI-dialect adapter posts to.
    ///
    /// An explicit `base_url` always wins, so a named provider can still be pointed at a proxy or
    /// a self-hosted gateway; otherwise the profile supplies it.
    pub fn resolve_base_url(&self) -> Result<String> {
        if let Some(url) = self
            .base_url
            .as_deref()
            .map(str::trim)
            .filter(|u| !u.is_empty())
        {
            return Ok(url.to_string());
        }
        if let Some(profile) = self.provider.profile() {
            return Ok(profile.base_url.to_string());
        }
        bail!(
            "provider = \"openai-compatible\" needs a `base_url` in \
             ~/.config/procyon/config.toml. For a known vendor, name it instead \
             (provider = \"deepseek\" | \"groq\" | \"openrouter\" | \"ollama\" | ...) \
             and the endpoint is preset."
        )
    }
}

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

    fn with_provider(provider: Provider) -> AppConfig {
        AppConfig {
            provider,
            ..Default::default()
        }
    }

    // The first suggestion is what a user who picks "ollama" and nothing else runs, and the whole
    // harness is tool calls — a model that cannot emit one turns every request into a polite
    // refusal that looks like a missing feature. 3B and 4B tags cannot; 8B is the floor.
    #[test]
    fn the_default_local_model_is_one_that_can_call_tools() {
        let default = Provider::Ollama.default_model();
        for too_small in [":1b", ":3b", ":4b", "llama3.2"] {
            assert!(
                !default.contains(too_small),
                "{} is too small to drive this harness reliably",
                default
            );
        }
    }

    // Naming a vendor is the whole point of the profiles: the user should not have to look up an
    // endpoint that never changes.
    #[test]
    fn a_named_provider_presets_its_endpoint() {
        assert_eq!(
            with_provider(Provider::Deepseek)
                .resolve_base_url()
                .unwrap(),
            "https://api.deepseek.com/v1"
        );
        assert_eq!(
            with_provider(Provider::Groq).resolve_base_url().unwrap(),
            "https://api.groq.com/openai/v1"
        );
    }

    // A named provider still has to be redirectable, or a proxy or self-hosted gateway forces the
    // user back to the generic variant and loses the credential preset with it.
    #[test]
    fn an_explicit_base_url_overrides_the_profile() {
        let config = AppConfig {
            provider: Provider::Deepseek,
            base_url: Some("https://my-gateway.internal/v1".to_string()),
            ..Default::default()
        };

        assert_eq!(
            config.resolve_base_url().unwrap(),
            "https://my-gateway.internal/v1"
        );
    }

    #[test]
    fn the_generic_variant_demands_a_base_url() {
        let err = with_provider(Provider::OpenAiCompatible)
            .resolve_base_url()
            .unwrap_err()
            .to_string();

        assert!(err.contains("base_url"), "{}", err);
        // The message has to name the easier path, not just reject the config.
        assert!(err.contains("deepseek"), "{}", err);
    }

    #[test]
    fn each_named_provider_defaults_to_its_own_credential_variable() {
        assert_eq!(
            with_provider(Provider::Openrouter).key_env_var(),
            "OPENROUTER_API_KEY"
        );
        assert_eq!(with_provider(Provider::Xai).key_env_var(), "XAI_API_KEY");
        assert_eq!(
            with_provider(Provider::Anthropic).key_env_var(),
            "ANTHROPIC_API_KEY"
        );
    }

    #[test]
    fn api_key_env_overrides_the_profile_default() {
        let config = AppConfig {
            provider: Provider::Deepseek,
            api_key_env: Some("MY_KEY".to_string()),
            ..Default::default()
        };

        assert_eq!(config.key_env_var(), "MY_KEY");
    }

    // A local server checks no credential, and demanding one would block the only way to run
    // Procyon with no provider account at all.
    #[test]
    fn a_local_provider_needs_no_credential() {
        assert_eq!(
            with_provider(Provider::Ollama).get_api_key().unwrap(),
            String::new()
        );
        assert_eq!(
            with_provider(Provider::LmStudio).get_api_key().unwrap(),
            String::new()
        );
    }

    #[test]
    fn a_remote_provider_without_a_credential_names_the_variable_to_set() {
        let err = with_provider(Provider::Groq)
            .get_api_key()
            .unwrap_err()
            .to_string();

        assert!(err.contains("GROQ_API_KEY"), "{}", err);
        // The error has to name the in-app fix, not just the environment-variable one.
        assert!(err.contains("/login"), "{}", err);
    }

    // `/login` writes here, so a stored key has to reach `get_api_key` before the config is asked
    // to fall back to an environment variable that may not exist at all.
    #[test]
    fn a_stored_credential_is_used_before_the_env_var() {
        let temp = tempfile::tempdir().unwrap();
        let mut store =
            crate::credentials::CredentialStore::load(temp.path().join("credentials.toml"))
                .unwrap();
        store.set("groq", "stored-key".to_string()).unwrap();

        let config = with_provider(Provider::Groq);
        assert_eq!(
            config.get_api_key_with_store(Some(&store)).unwrap(),
            "stored-key"
        );
    }

    #[test]
    fn no_stored_credential_falls_through_to_the_env_var_path() {
        let temp = tempfile::tempdir().unwrap();
        let store = crate::credentials::CredentialStore::load(temp.path().join("credentials.toml"))
            .unwrap();

        let err = with_provider(Provider::Groq)
            .get_api_key_with_store(Some(&store))
            .unwrap_err()
            .to_string();
        assert!(err.contains("GROQ_API_KEY"), "{}", err);
    }

    fn parse(theme_value: &str) -> Result<AppConfig> {
        let toml_str = format!(
            "default_network = \"testnet\"\ntheme = \"{}\"\ndefault_model = \"m\"\n",
            theme_value
        );
        Ok(toml::from_str(&toml_str)?)
    }

    #[test]
    fn deserializes_lowercase_theme_from_config_file() {
        assert_eq!(parse("dark").unwrap().theme, Theme::Dark);
        assert_eq!(parse("light").unwrap().theme, Theme::Light);
    }

    #[test]
    fn a_config_may_set_only_some_fields() {
        let cfg: AppConfig = toml::from_str("theme = \"light\"\n").expect("partial config");
        assert_eq!(cfg.theme, Theme::Light);
        assert_eq!(cfg.default_network, AppConfig::default().default_network);
    }

    #[test]
    fn an_empty_config_is_all_defaults() {
        let cfg: AppConfig = toml::from_str("").expect("empty config");
        assert_eq!(cfg.theme, AppConfig::default().theme);
    }

    #[test]
    fn rejects_unknown_theme() {
        assert!(parse("solarized").is_err());
    }

    #[test]
    fn serde_agrees_with_display() {
        for theme in [Theme::Dark, Theme::Light] {
            let written = toml::to_string(&AppConfig {
                theme: theme.clone(),
                ..AppConfig::default()
            })
            .unwrap();
            assert!(
                written.contains(&format!("theme = \"{}\"", theme)),
                "Display spelling of {:?} is not what serde wrote: {}",
                theme,
                written
            );
        }
    }

    #[test]
    fn serialized_theme_round_trips() {
        let written = toml::to_string(&AppConfig::default()).unwrap();
        let reloaded: AppConfig = toml::from_str(&written).unwrap();
        assert_eq!(reloaded.theme, AppConfig::default().theme);
    }

    #[test]
    fn from_str_matches_serde() {
        for spelling in ["dark", "light"] {
            assert_eq!(
                Theme::from_str(spelling).unwrap(),
                parse(spelling).unwrap().theme
            );
        }
    }
}