churn-v2/crates/churn/src/agent/config.rs

use std::collections::BTreeMap;

use anyhow::Context;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Clone)]
pub struct AgentConfig {
    pub agent_id: String,
    pub discovery: String,

    pub labels: BTreeMap<String, String>,
}

impl AgentConfig {
    pub async fn new() -> anyhow::Result<Self> {
        let config = ConfigFile::load().await?;

        Ok(Self {
            agent_id: config.agent_id,
            discovery: config.discovery,
            labels: config.labels.unwrap_or_default(),
        })
    }
}

#[derive(Serialize, Deserialize)]
struct ConfigFile {
    agent_id: String,
    discovery: String,

    labels: Option<BTreeMap<String, String>>,
}

impl ConfigFile {
    pub async fn load() -> anyhow::Result<Self> {
        let directory = dirs::data_dir()
            .ok_or(anyhow::anyhow!("failed to get data dir"))?
            .join("io.kjuulh.churn-agent")
            .join("churn-agent.toml");

        if !directory.exists() {
            anyhow::bail!(
                "No churn agent file was setup, run `churn agent setup` to setup the defaults"
            )
        }

        let contents = tokio::fs::read_to_string(&directory).await?;

        toml::from_str(&contents).context("failed to parse the contents of the churn agent config")
    }

    pub async fn write_default(
        discovery: impl Into<String>,
        force: bool,
        labels: impl Into<BTreeMap<String, String>>,
    ) -> anyhow::Result<Self> {
        let s = Self {
            agent_id: Uuid::new_v4().to_string(),
            discovery: discovery.into(),
            labels: Some(labels.into()),
        };

        let directory = dirs::data_dir()
            .ok_or(anyhow::anyhow!("failed to get data dir"))?
            .join("io.kjuulh.churn-agent")
            .join("churn-agent.toml");

        if let Some(parent) = directory.parent() {
            tokio::fs::create_dir_all(&parent).await?;
        }

        if !force && directory.exists() {
            anyhow::bail!("config file already exists, consider moving it to a backup before trying again: {}", directory.display());
        }

        let contents = toml::to_string_pretty(&s)
            .context("failed to convert default implementation to string")?;

        tokio::fs::write(directory, contents.as_bytes())
            .await
            .context("failed to write to agent file")?;

        Ok(s)
    }
}

pub async fn setup_config(
    discovery: impl Into<String>,
    force: bool,
    labels: impl Into<BTreeMap<String, String>>,
) -> anyhow::Result<()> {
    ConfigFile::write_default(discovery, force, labels).await?;

    Ok(())
}