monothink 0.2.0

/*
 * This file is part of monothink.
 *
 * Wavething is free software: you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free Software Foundation,
 * either version 3 of the License, or (at your option) any later version.
 *
 * Wavething is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with Wavething.
 * If not, see <https://www.gnu.org/licenses/>.
 */

use anyhow::{Context, Result};
use directories::ProjectDirs;
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, fs::File, io::Read, path::Path};

#[derive(Debug, Clone, Default, Deserialize)]
pub struct Auth {
    pub endpoint: BTreeMap<String, EndpointSecrets>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct EndpointSecrets {
    pub api_token: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Config {
    pub general: General,
    pub endpoint: BTreeMap<String, Endpoint>,
    #[serde(default)]
    pub model: BTreeMap<String, Model>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct General {
    pub primary_endpoint: String,
    #[serde(default = "default_max_tokens_per_iteration")]
    pub max_tokens: u64,
    #[serde(default = "default_max_iterations_per_request")]
    pub max_iterations_per_request: usize,
    #[serde(default = "default_max_dir_list_results")]
    pub max_list_dir_results: usize,
    #[serde(default = "default_max_search_text_results")]
    pub max_search_text_results: usize,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Endpoint {
    pub base_url: String,
    pub default_model: String,
    #[serde(skip)]
    pub secrets: Option<EndpointSecrets>,
    #[serde(default)]
    pub models: Vec<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Model {
    #[serde(default = "default_context_window")]
    pub context_window: u64,
}

fn default_context_window() -> u64 {
    32000
}

fn default_max_tokens_per_iteration() -> u64 {
    8192
}

fn default_max_iterations_per_request() -> usize {
    40
}

fn default_max_dir_list_results() -> usize {
    100
}

fn default_max_search_text_results() -> usize {
    100
}

// @todo overall, make error nicer using with_message (e.g. which file was not found or which value was expected?)
// load config and secrets
pub fn load_config() -> Result<Config> {
    // load config from default path
    if let Some(proj_dirs) = ProjectDirs::from("org", "codemonk", "monothink") {
        let config_dir = proj_dirs.config_dir().to_path_buf();

        let mut config_file_dir = config_dir.clone();
        config_file_dir.push("config.toml");
        let config_file_path = config_file_dir.as_path();

        let mut auth_file_dir = config_dir.clone();
        auth_file_dir.push("auth.toml");
        let auth_file_path = auth_file_dir.as_path();

        let mut config = parse_config(config_file_path).with_context(|| {
            format!(
                "Error while parsing config file '{}'.",
                config_file_path.to_string_lossy()
            )
        })?;
        let auth = parse_auth(auth_file_path).with_context(|| {
            format!(
                "Error while parsing auth file '{}'.",
                auth_file_path.to_string_lossy()
            )
        })?;
        merge_config_with_auth(&mut config, &auth);

        Ok(config)
    } else {
        panic!("Could not retrieve a valid home directory from the operating system.");
    }
}

fn merge_config_with_auth(config: &mut Config, auth: &Auth) {
    for (key, val) in &auth.endpoint {
        if let Some(endpoint) = config.endpoint.get_mut(key) {
            endpoint.secrets = Some(val.to_owned());
        } else {
            // @todo if a key has no matching endpoint, print a warning
        }
        // @todo if a key is missing for a endpoint, print a warning
        // for this iterate over all endpoints in check if secrets is none after merge
    }
}

fn parse_auth(path: &Path) -> Result<Auth> {
    let mut file = File::open(path)?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)?;

    let auth: Auth = toml::from_str(&contents)?;
    // @todo if it fails, print an example
    Ok(auth)
}

fn parse_config(path: &Path) -> Result<Config> {
    let mut file = File::open(path)?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)?;

    let config: Config = toml::from_str(&contents)?;
    // @todo if it fails, print an example
    Ok(config)
}