overbrainer 0.10.0

use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};

use config::{Config, Environment, File, FileFormat};
use serde_json::Value;

use super::{Settings, validate};

/// Name of the project-level configuration file, relative to the project directory.
pub const CONFIG_FILE: &str = "overbrainer.toml";
/// Prefix required on every environment variable read into the configuration.
pub const ENV_PREFIX: &str = "OVERBRAINER";
/// The `OVERBRAINER_*` variables read outside the configuration, which [`load`]
/// skips: the terminal UI's own (`OVERBRAINER_TUI_COLOR`, `OVERBRAINER_TUI_MOTION`),
/// the test suite's (`OVERBRAINER_TEST_SSH_HOST`, `OVERBRAINER_TEST_SSH_CONFIG`) and
/// the update check's switch and the SSH client's ([`super::SSH_CLIENT_ENV`], read by
/// [`ssh_client_env`]). A name ending in `_` is a prefix, any other is exact.
const NOT_CONFIG_VARIABLES: [&str; 4] = [
    "OVERBRAINER_TUI_",
    "OVERBRAINER_TEST_",
    crate::update::NO_UPDATE_CHECK_ENV,
    super::SSH_CLIENT_ENV,
];

/// Where [`load`] reads `OVERBRAINER_*` environment variable overrides from.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvSource {
    /// Read from the current process environment. This is what a running binary uses.
    Process,
    /// Use these key-value pairs instead of the process environment. Tests use this
    /// to inject environment variables without calling `std::env::set_var`.
    Vars(Vec<(String, String)>),
}

/// Where a command reads its settings: `overbrainer.toml` on disk, or a text
/// already read and validated (the configuration the TUI keeps), layered with
/// the variables of `env`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Source {
    /// The text of `overbrainer.toml`; `None` reads the file.
    pub text: Option<String>,
    /// Where the `OVERBRAINER_*` variables come from.
    pub env: EnvSource,
}

impl Source {
    /// The settings of `project_dir` from this source.
    ///
    /// # Errors
    ///
    /// Returns what [`load`] returns without a text, else what [`load_str`]
    /// returns.
    pub fn load(&self, project_dir: &Path) -> Result<Settings, ConfigError> {
        match &self.text {
            Some(text) => load_str(text, self.env.clone()),
            None => load(project_dir, self.env.clone()),
        }
    }
}

impl From<EnvSource> for Source {
    /// The file on disk, with `env`.
    fn from(env: EnvSource) -> Self {
        Self { text: None, env }
    }
}

/// Everything that can go wrong while loading configuration.
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
    /// The configuration file could not be read. The I/O error is part of the
    /// message and not a source, so a chain of causes (`{:#}`) names it once.
    #[error("cannot read {}: {error}", path.display())]
    Read {
        /// Path to the file that could not be read.
        path: PathBuf,
        /// Underlying I/O error.
        error: std::io::Error,
    },
    /// The file or environment could not be parsed into `Settings`. The message names
    /// the key and the expected type, never the offending value, which may be a secret
    /// set in the wrong variable.
    #[error("invalid configuration: {0}")]
    Parse(String),
    /// The parsed settings failed semantic validation.
    #[error("invalid configuration:\n  {}", .0.join("\n  "))]
    Invalid(Vec<String>),
}

impl From<config::ConfigError> for ConfigError {
    fn from(source: config::ConfigError) -> Self {
        Self::Parse(describe(&source))
    }
}

/// Describes a `config` error without the offending value.
///
/// Type errors keep their key and expected type. Custom messages from serde quote the
/// value (`invalid value: string "..."`, `unknown variant ...`), so only the part from
/// `expected` onward is kept. Unknown field names are keys, not values, and are kept.
/// TOML syntax errors are reported by location only, never with source snippets.
fn describe(error: &config::ConfigError) -> String {
    match error {
        config::ConfigError::Type { key, expected, .. } => {
            let key = key.as_deref().unwrap_or("a value");
            format!("`{key}` has the wrong type, expected {expected}")
        },
        config::ConfigError::At { error, key, .. } => match key {
            Some(key) => format!("`{key}`: {}", describe(error)),
            None => describe(error),
        },
        config::ConfigError::Message(message) => redact(message),
        config::ConfigError::FileParse { .. } => redact_toml_error(error),
        _ => "invalid configuration".to_string(),
    }
}

/// Keeps only the location of a TOML parse error ("at line N, column M"), never the
/// source snippet that follows it on the next lines.
fn redact_toml_error(error: &config::ConfigError) -> String {
    let message = error.to_string();
    let first_line = message.lines().next().unwrap_or("");
    match first_line.find("at line ") {
        Some(pos) => format!(
            "TOML syntax error in overbrainer.toml {}",
            first_line[pos..].trim_end()
        ),
        None => "TOML syntax error in overbrainer.toml".to_string(),
    }
}

/// Keeps what a serde message says was expected, dropping the value it quotes.
fn redact(message: &str) -> String {
    if message.starts_with("unknown field") || message.starts_with("missing field") {
        return message.to_string();
    }
    if let Some(index) = message.find("expected") {
        return format!("invalid value, {}", &message[index..]);
    }
    if let Some((head, _)) = message.split_once(" does not have variant constructor") {
        return format!("{head}: unknown variant");
    }
    "invalid value".to_string()
}

/// Loads `<project_dir>/overbrainer.toml` layered with `OVERBRAINER_*` variables.
///
/// `env: EnvSource::Process` reads the process environment; tests pass
/// `EnvSource::Vars` with an explicit list of key-value pairs.
///
/// # Errors
///
/// Returns [`ConfigError::Read`] when the file cannot be read, [`ConfigError::Parse`]
/// when the file or environment cannot be parsed into [`Settings`], and
/// [`ConfigError::Invalid`] when the parsed settings fail semantic validation or
/// when an env-only key is set in the file.
pub fn load(project_dir: &Path, env: EnvSource) -> Result<Settings, ConfigError> {
    let path = project_dir.join(CONFIG_FILE);
    let content = std::fs::read_to_string(&path).map_err(|error| ConfigError::Read {
        path: path.clone(),
        error,
    })?;
    load_str(&content, env)
}

/// Validates a configuration text (the would-be contents of `overbrainer.toml`) layered
/// with `OVERBRAINER_*` variables from `env`, without reading any file.
///
/// This is what [`load`] does after reading the file; a caller that already has the
/// text (an editor buffer, a value about to be written) can validate it directly.
///
/// # Errors
///
/// Returns [`ConfigError::Parse`] when `content` or the environment cannot be parsed
/// into [`Settings`], and [`ConfigError::Invalid`] when the parsed settings fail
/// semantic validation or when an env-only key is set in `content`.
pub fn load_str(content: &str, env: EnvSource) -> Result<Settings, ConfigError> {
    let env: config::Map<String, String> = env_pairs(env)
        .into_iter()
        .filter(|(key, _)| !is_not_config(key))
        .collect();

    let file_only = Config::builder()
        .add_source(File::from_str(content, FileFormat::Toml))
        .build()?;
    let mut problems = validate::env_only_in_file(&file_only);

    let mut settings: Settings = Config::builder()
        .add_source(File::from_str(content, FileFormat::Toml))
        .add_source(
            Environment::with_prefix(ENV_PREFIX)
                .prefix_separator("_")
                .separator("__")
                .try_parsing(false)
                .source(Some(env)),
        )
        .build()?
        .try_deserialize()?;

    if let Some(training) = settings.training.as_mut() {
        let file_extra: BTreeMap<String, Value> =
            file_only.get("training.axolotl_extra").unwrap_or_default();
        coerce_env_scalars(&mut training.axolotl_extra, &file_extra);
    }

    problems.extend(validate::check(&settings));
    if problems.is_empty() {
        Ok(settings)
    } else {
        Err(ConfigError::Invalid(problems))
    }
}

/// The dotted, lower-case configuration keys `env` sets through `OVERBRAINER_*`
/// variables, with the same prefix and separators [`load`] and [`load_str`] use.
///
/// `OVERBRAINER_PROVIDERS__OPENROUTER__API_KEY` gives `providers.openrouter.api_key`;
/// `OVERBRAINER_LOG` gives `log`. The variables read outside the configuration
/// (`OVERBRAINER_TUI_*`, `OVERBRAINER_TEST_*`, `OVERBRAINER_NO_UPDATE_CHECK`) are not
/// configuration keys ([`load_str`] skips them too) and are left out.
///
/// The prefix and those names are matched case-insensitively, on a lower-cased
/// copy of each key: this mirrors what `config::Environment` itself does inside
/// [`load_str`] (it lower-cases every key before comparing it to its own lower-cased
/// prefix pattern), so a variable such as `overbrainer_log` or `Overbrainer_Log`,
/// which `load_str` accepts, is reported here too.
#[must_use]
pub fn env_keys(env: &EnvSource) -> BTreeSet<String> {
    let prefix = format!("{ENV_PREFIX}_").to_lowercase();
    env_pairs(env.clone())
        .into_iter()
        .filter(|(key, _)| !is_not_config(key))
        .map(|(key, _)| key.to_lowercase())
        .filter_map(|key| {
            let rest = key.strip_prefix(&prefix)?;
            (!rest.is_empty()).then(|| rest.replace("__", "."))
        })
        .collect()
}

/// The value of `OVERBRAINER_SSH_CLIENT` in `env`, for [`super::effective_client`].
#[must_use]
pub fn ssh_client_env(env: &EnvSource) -> Option<String> {
    env_pairs(env.clone())
        .into_iter()
        .find(|(key, _)| key.eq_ignore_ascii_case(super::SSH_CLIENT_ENV))
        .map(|(_, value)| value)
}

/// The client a target set to `configured` uses in this process:
/// [`super::effective_client`] with `OVERBRAINER_SSH_CLIENT` from the process
/// environment.
///
/// # Errors
///
/// Returns the message of [`super::effective_client`], which names the
/// variable when its value is invalid.
pub fn process_client(configured: super::SshClient) -> Result<super::SshClient, String> {
    super::effective_client(configured, ssh_client_env(&EnvSource::Process).as_deref())
}

/// Whether `key` is one of [`NOT_CONFIG_VARIABLES`] once lower-cased, as
/// `config::Environment` lower-cases keys before matching.
fn is_not_config(key: &str) -> bool {
    let key = key.to_lowercase();
    NOT_CONFIG_VARIABLES.iter().any(|name| {
        let name = name.to_lowercase();
        if name.ends_with('_') {
            key.starts_with(&name)
        } else {
            key == name
        }
    })
}

/// The key-value pairs `env` provides: the process environment, or an explicit list.
fn env_pairs(env: EnvSource) -> Vec<(String, String)> {
    match env {
        EnvSource::Process => std::env::vars_os()
            .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?)))
            .collect(),
        EnvSource::Vars(pairs) => pairs,
    }
}

/// Gives env-provided values of `training.axolotl_extra` a type.
///
/// The `Environment` source above is built with `try_parsing(false)`, the config-wide
/// default that keeps every env value a string (so a secret such as `0123` is never
/// silently type-parsed). This function is where `axolotl_extra` alone deliberately
/// relaxes that default: every value set through
/// `OVERBRAINER_TRAINING__AXOLOTL_EXTRA__*` arrives as a string, and a string that is
/// not the file's own value for the same key came from env and becomes a boolean
/// (`true`, `false`), a canonical integer, or a finite float when it parses as one
/// (see [`scalar`]); any other text stays a string. Values from `overbrainer.toml`
/// keep their TOML type. A value that must stay a string, such as a revision with
/// leading zeros (`"0123"`) or a version-like `"3.14"`, belongs in the file; env has
/// no way to know a key means "always a string", so an override of that key with
/// different text is still coerced by this function.
fn coerce_env_scalars(extra: &mut BTreeMap<String, Value>, file: &BTreeMap<String, Value>) {
    for (key, value) in extra.iter_mut() {
        coerce(value, file.get(key));
    }
}

fn coerce(value: &mut Value, file: Option<&Value>) {
    if let Value::Object(map) = value {
        for (key, inner) in map.iter_mut() {
            coerce(inner, file.and_then(|file| file.get(key)));
        }
        return;
    }
    let parsed = match &*value {
        Value::String(text) if file.and_then(Value::as_str) != Some(text.as_str()) => scalar(text),
        _ => None,
    };
    if let Some(parsed) = parsed {
        *value = parsed;
    }
}

/// `text` as a boolean, an integer or a finite float, when it is one.
///
/// An integer is only recognized when `text` is its own canonical rendering (no
/// leading zero, no leading `+`), so a padded value such as a revision (`"0123"`)
/// is not silently read as `123`. `text` is first tried as an `i64`, then, for a
/// positive value that overflows it, as a `u64`, so a value up to [`u64::MAX`]
/// keeps its exact precision. When `text` reads as an integer (an optional sign
/// followed only by digits) but is not the canonical rendering of an `i64` or
/// `u64`, either because it overflows both or because of a leading zero or `+`, it
/// is left as a string: it is never tried as a float, which would silently lose
/// precision for a value outside the `i64`/`u64` range. Floats have no such check:
/// any other text that parses as a finite `f64` is accepted.
fn scalar(text: &str) -> Option<Value> {
    match text {
        "true" => return Some(Value::Bool(true)),
        "false" => return Some(Value::Bool(false)),
        _ => {},
    }
    if let Some(integer) = canonical_integer(text) {
        return Some(integer);
    }
    if looks_like_integer(text) {
        return None;
    }
    text.parse::<f64>()
        .ok()
        .and_then(serde_json::Number::from_f64)
        .map(Value::Number)
}

/// `text` as an `i64`, or, for a positive value that overflows `i64`, as a `u64`,
/// but only when `text` is that integer's own canonical rendering (no leading
/// zero, no leading `+`).
fn canonical_integer(text: &str) -> Option<Value> {
    if let Ok(integer) = text.parse::<i64>() {
        return (integer.to_string() == text).then_some(Value::from(integer));
    }
    if let Ok(integer) = text.parse::<u64>() {
        return (integer.to_string() == text).then_some(Value::from(integer));
    }
    None
}

/// True when `text` reads as an integer: an optional leading `+` or `-` followed by
/// one or more ASCII digits. This is broader than [`canonical_integer`], which also
/// requires `text` to be that integer's canonical rendering and to fit in an `i64`
/// or `u64`.
fn looks_like_integer(text: &str) -> bool {
    let digits = text.strip_prefix(['+', '-']).unwrap_or(text);
    !digits.is_empty() && digits.bytes().all(|byte| byte.is_ascii_digit())
}

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

    #[test]
    fn scalars_parse_only_plain_values() {
        assert_eq!(scalar("true"), Some(Value::Bool(true)));
        assert_eq!(scalar("10"), Some(Value::from(10)));
        assert_eq!(scalar("0.05"), Some(Value::from(0.05)));
        assert_eq!(scalar("nan"), None);
        assert_eq!(scalar("inf"), None);
        assert_eq!(scalar("qwen3"), None);
        assert_eq!(scalar("True"), None);
    }

    #[test]
    fn non_canonical_integers_stay_strings() {
        assert_eq!(scalar("0123"), None, "leading zero");
        assert_eq!(scalar("+5"), None, "leading plus");
        assert_eq!(scalar("007"), None, "leading zeros");
    }

    #[test]
    fn canonical_integers_and_floats_still_parse() {
        assert_eq!(scalar("0"), Some(Value::from(0)));
        assert_eq!(scalar("-3"), Some(Value::from(-3)));
        assert_eq!(scalar("1e-7"), Some(Value::from(1e-7)));
    }

    #[test]
    fn i64_max_still_parses_as_an_integer() {
        assert_eq!(scalar("9223372036854775807"), Some(Value::from(i64::MAX)));
    }

    #[test]
    fn i64_max_plus_one_becomes_a_u64_not_a_float() {
        assert_eq!(
            scalar("9223372036854775808"),
            Some(Value::from(9_223_372_036_854_775_808_u64)),
            "one past i64::MAX must not lose precision as an f64"
        );
    }

    #[test]
    fn u64_max_still_parses_as_an_integer() {
        assert_eq!(scalar("18446744073709551615"), Some(Value::from(u64::MAX)));
    }

    #[test]
    fn u64_max_plus_one_stays_a_string() {
        assert_eq!(
            scalar("18446744073709551616"),
            None,
            "an integer past u64::MAX must stay a string, not become an imprecise float"
        );
    }

    #[test]
    fn a_negative_integer_below_i64_min_stays_a_string() {
        assert_eq!(
            scalar("-9223372036854775809"),
            None,
            "a negative integer past i64::MIN must stay a string, not become an imprecise float"
        );
    }

    fn vars(pairs: &[(&str, &str)]) -> EnvSource {
        EnvSource::Vars(
            pairs
                .iter()
                .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
                .collect(),
        )
    }

    #[test]
    fn env_keys_maps_double_underscores_to_dots_and_lower_cases() {
        let env = vars(&[("OVERBRAINER_PROVIDERS__OPENROUTER__API_KEY", "sk-something")]);
        assert_eq!(
            env_keys(&env),
            BTreeSet::from(["providers.openrouter.api_key".to_string()])
        );
    }

    #[test]
    fn env_keys_maps_log_to_a_bare_key() {
        let env = vars(&[("OVERBRAINER_LOG", "debug")]);
        assert_eq!(env_keys(&env), BTreeSet::from(["log".to_string()]));
    }

    /// `OVERBRAINER_SSH_CLIENT` is read by [`ssh_client_env`], not loaded as a key.
    #[test]
    fn the_ssh_client_variable_is_read_apart_from_the_configuration() {
        let env = vars(&[("OVERBRAINER_SSH_CLIENT", "builtin")]);
        assert_eq!(env_keys(&env), BTreeSet::new());
        assert_eq!(ssh_client_env(&env), Some("builtin".to_string()));
        assert_eq!(ssh_client_env(&vars(&[])), None);
    }

    /// A lower-case `overbrainer_ssh_client` is read like the upper-case name.
    #[test]
    fn the_ssh_client_variable_is_read_in_any_case() {
        let env = vars(&[("overbrainer_ssh_client", "builtin")]);
        assert_eq!(ssh_client_env(&env), Some("builtin".to_string()));
    }

    #[test]
    fn env_keys_skips_tui_variables_and_unrelated_prefixes() {
        let env = vars(&[
            ("OVERBRAINER_TUI_COLOR", "256"),
            ("OVERBRAINER_TUI_MOTION", "off"),
            ("PATH", "/usr/bin"),
            ("OVERBRAINERX", "not us"),
        ]);
        assert!(env_keys(&env).is_empty());
    }

    #[test]
    fn env_keys_skips_the_variables_read_outside_the_configuration() {
        let env = vars(&[
            ("OVERBRAINER_NO_UPDATE_CHECK", "1"),
            ("overbrainer_no_update_check", "1"),
            ("OVERBRAINER_TEST_SSH_HOST", "overbrainer-test"),
        ]);
        assert!(env_keys(&env).is_empty());
    }

    #[test]
    fn env_keys_skips_the_bare_prefix_with_nothing_after_it() {
        let env = vars(&[("OVERBRAINER_", "")]);
        assert!(env_keys(&env).is_empty());
    }

    #[test]
    fn env_keys_matches_a_lower_case_variable_name() {
        // `config::Environment` lower-cases every key before matching its prefix, so
        // `load_str` accepts `overbrainer_log` exactly as it accepts `OVERBRAINER_LOG`.
        let env = vars(&[("overbrainer_log", "debug")]);
        assert_eq!(env_keys(&env), BTreeSet::from(["log".to_string()]));
    }

    #[test]
    fn env_keys_matches_a_mixed_case_variable_name() {
        let env = vars(&[("Overbrainer_Providers__Openrouter__Api_Key", "sk-something")]);
        assert_eq!(
            env_keys(&env),
            BTreeSet::from(["providers.openrouter.api_key".to_string()])
        );
    }

    #[test]
    fn env_keys_skips_a_lower_case_tui_variable_case_insensitively() {
        let env = vars(&[("overbrainer_tui_color", "256")]);
        assert!(env_keys(&env).is_empty());
    }
}