robotnik/src/setup.rs

use clap::Parser;
use color_eyre::{
    Result,
    eyre::WrapErr,
};
use config::Config;
use directories::ProjectDirs;
use std::path::PathBuf;
use tracing::{
    info,
    instrument,
};

// TODO: use [clap(long, short, help_heading = Some(section))]
#[derive(Clone, Debug, Parser)]
#[command(about, version)]
pub(crate) struct Args {
    #[arg(short, long)]
    /// API Key for the LLM in use.
    pub(crate) api_key: Option<String>,

    #[arg(short, long, default_value = "https://api.openai.com")]
    /// Base URL for the LLM API to use.
    pub(crate) base_url: Option<String>,

    /// Directory to use for chroot (recommended).
    #[arg(long)]
    pub(crate) chroot_dir: Option<String>,

    /// Root directory for file based command structure.
    #[arg(long)]
    pub(crate) command_dir: Option<String>,

    #[arg(long)]
    /// Instructions to the model on how to behave.
    pub(crate) instruct: Option<String>,

    #[arg(long)]
    pub(crate) model: Option<String>,

    #[arg(long)]
    /// List of IRC channels to join.
    pub(crate) channels: Option<Vec<String>>,

    #[arg(short, long)]
    /// Custom configuration file location if need be.
    pub(crate) config_file: Option<PathBuf>,

    #[arg(short, long, default_value = "irc.libera.chat")]
    /// IRC server.
    pub(crate) server: Option<String>,

    #[arg(short, long, default_value = "6697")]
    /// Port of the IRC server.
    pub(crate) port: Option<String>,

    #[arg(long)]
    /// IRC Nickname.
    pub(crate) nickname: Option<String>,

    #[arg(long)]
    /// IRC Nick Password
    pub(crate) nick_password: Option<String>,

    #[arg(long)]
    /// IRC Username
    pub(crate) username: Option<String>,

    #[arg(long)]
    /// Whether or not to use TLS when connecting to the IRC server.
    pub(crate) use_tls: Option<bool>,
}

pub(crate) struct Setup {
    pub(crate) config: Config,
}

#[instrument]
pub async fn init() -> Result<Setup> {
    // Get arguments.  These overrule configuration file, and environment
    // variables if applicable.
    let args = Args::parse();

    // Use default config location unless specified.
    let config_location: PathBuf = if let Some(ref path) = args.config_file {
        path.to_owned()
    } else {
        ProjectDirs::from("", "", env!("CARGO_PKG_NAME"))
            .unwrap()
            .config_dir()
            .to_owned()
            .join(r"config.toml")
    };

    info!("Starting.");

    let settings = Config::builder()
        .add_source(config::File::with_name(&config_location.to_string_lossy()).required(false))
        .add_source(config::Environment::with_prefix("BOT"))
        // Doing all of these overrides provides a unified access point for options,
        // but a derive macro could do this a bit better if this becomes too large.
        .set_override_option("api-key", args.api_key.clone())?
        .set_override_option("base-url", args.base_url.clone())?
        .set_override_option("chroot-dir", args.chroot_dir.clone())?
        .set_override_option("command-path", args.command_dir.clone())?
        .set_override_option("model", args.model.clone())?
        .set_override_option("instruct", args.instruct.clone())?
        .set_override_option("channels", args.channels.clone())?
        .set_override_option("server", args.server.clone())?
        .set_override_option("port", args.port.clone())? // FIXME: Make this a default here not in clap.
        .set_override_option("nickname", args.nickname.clone())?
        .set_override_option("username", args.username.clone())?
        .set_override_option("use_tls", args.use_tls)?
        .build()
        .wrap_err("Couldn't read configuration settings.")?;

    Ok(Setup { config: settings })
}