iocaine 3.3.0

The deadliest poison known to AI
// SPDX-FileCopyrightText: Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

use anyhow::anyhow;
use figment::{
    Figment,
    providers::{Format, Json, Toml, Yaml},
};
use figment_file_provider_adapter::{FileAdapter, FileAdapterWithRestrictions};
use glob::glob;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use tokio_listener::ListenerAddress;

use crate::{
    bruce::Bruce,
    checkpoint_charlie::CheckpointCharlie,
    morgue::freezer::{self, Freezer},
    tenx_programmer::TenXProgrammer,
};

use iocaine_powder::{acab::ACAB, sex_dungeon::Language};

mod cuddles;
use cuddles::Cuddles;

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    #[serde(default)]
    pub initial_seed: String,
    #[serde(default)]
    pub state_directory: PathBuf,
    pub instance_id: Option<String>,
    #[serde(default, rename = "server")]
    pub servers: HashMap<String, Server>,
    #[serde(default, rename = "handler")]
    pub handlers: HashMap<String, Handler>,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub struct Server {
    pub bind: ListenerAddress,
    pub unix_socket_access: Option<tokio_listener::UnixChmodVariant>,

    #[serde(flatten)]
    pub variant: Variant,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(tag = "mode")]
pub enum Variant {
    #[serde(rename = "http")]
    Http(RequestHandler),
    #[serde(rename = "haproxy-spoa")]
    HaproxySPOA(RequestHandler),
    #[serde(rename = "prometheus")]
    Prometheus(PrometheusConfig),
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub struct RequestHandler {
    pub initial_seed: Option<String>,
    #[serde(default, rename = "use")]
    pub uses: Uses,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Handler {
    pub path: Option<PathBuf>,
    #[serde(default)]
    pub language: Language,
    pub compiler: Option<PathBuf>,
    pub config: Option<figment::value::Dict>,
}

#[derive(Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub struct Uses {
    pub metrics: Option<String>,
    pub handler_from: String,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub struct PrometheusConfig {
    pub persist_path: Option<PathBuf>,
    #[serde(default = "PrometheusConfig::default_persist_interval")]
    pub persist_interval: String,
}

impl PrometheusConfig {
    fn default_persist_interval() -> String {
        String::from("1h")
    }
}

fn file_wrap<T: figment::Provider + 'static>(provider: T) -> FileAdapterWithRestrictions {
    FileAdapter::wrap(provider)
        .with_suffix("-file")
        .only(&["initial-seed"])
}

impl Config {
    fn init_state(&self) -> anyhow::Result<ACAB> {
        let state_path = self.state_directory.join("state.json");
        let mut acab = ACAB::load(&state_path).map_err(|e| anyhow::anyhow!("{e:#?}"))?;
        if let Some(id) = &self.instance_id {
            id.clone_into(&mut acab.state.instance_id);
        }
        acab.save().map_err(|e| anyhow::anyhow!("{e:#?}"))?;

        Ok(acab)
    }

    pub fn load(paths: &[String]) -> Result<Self, anyhow::Error> {
        let mut config = Figment::new().merge(file_wrap(Cuddles::default()));
        let mut config_files = Vec::new();

        for path in paths {
            if Path::new(path).is_dir() {
                let mut paths = Vec::new();
                for format in ["toml", "json", "yaml", "yml", "kdl"] {
                    let config_d = format!("{path}/**/*.{format}");
                    let mut files: Vec<_> = glob(&config_d)?.filter_map(Result::ok).collect();
                    paths.append(&mut files);
                }
                paths.sort();
                config_files.append(&mut paths);
            } else if Path::new(path).try_exists()? {
                config_files = Vec::from([PathBuf::from(path)]);
                config = Figment::new();
            } else {
                tracing::warn!({ config_file = path }, "configuration path does not exist");
            }
        }

        for file in &config_files {
            let Some(ext) = file.extension() else {
                anyhow::bail!(
                    "Unrecognized configuration file format for file: {}",
                    file.display()
                );
            };
            if ext == "toml" {
                config = config.merge(file_wrap(Toml::file(file)));
            } else if ext == "json" {
                config = config.merge(file_wrap(Json::file(file)));
            } else if ext == "yaml" || ext == "yml" {
                config = config.merge(file_wrap(Yaml::file(file)));
            } else if ext == "kdl" {
                config = config.merge(file_wrap(Cuddles::file(file)));
            } else {
                anyhow::bail!(
                    "Unrecognised configuration file format for file: {}",
                    file.display()
                );
            }

            tracing::debug!(
                { config_file = file.display().to_string() },
                "loading configuration"
            );
        }

        config
            .extract()
            .map_err(|e| anyhow!("Failed to load configuration: {e}"))
    }
}

impl TryFrom<Config> for Freezer {
    type Error = anyhow::Error;

    fn try_from(config: Config) -> Result<Self, Self::Error> {
        let mut servers = HashMap::new();

        // Pull out the metrics servers first
        for (name, server) in &config.servers {
            let Variant::Prometheus(cfg) = &server.variant else {
                continue;
            };

            let persist_interval = duration_str::parse(&cfg.persist_interval)
                .map_err(|e| anyhow::anyhow!("failed to parse persist-interval:\n{e}"))?;

            let srv = freezer::Server {
                bind: server.bind.clone(),
                unix_socket_access: server.unix_socket_access,
                variant: TenXProgrammer::new(cfg.persist_path.as_ref(), persist_interval)
                    .map_err(|e| anyhow::anyhow!("{e}"))?
                    .into(),
            };
            servers.insert(name.to_owned(), srv);
        }

        let find_prometheus_server = |servers: &HashMap<String, freezer::Server>,
                                      server_name: &str,
                                      metrics_name: &Option<String>|
         -> anyhow::Result<TenXProgrammer> {
            let metrics = match metrics_name {
                Some(metrics_name) => {
                    let Some(metrics_server) = servers.get(metrics_name) else {
                        anyhow::bail!("[server.{server_name}]: {metrics_name} not found");
                    };
                    let freezer::Variant::Prometheus(metrics_server) = &metrics_server.variant
                    else {
                        anyhow::bail!(
                            "[server.{server_name}]: {metrics_name} is not a prometheus server"
                        );
                    };
                    metrics_server.clone()
                }
                None => TenXProgrammer::default(),
            };
            Ok(metrics)
        };

        let find_handler = |name: &str| -> anyhow::Result<Handler> {
            let Some(handler) = config.handlers.get(name) else {
                anyhow::bail!("No handler declared with name: {name}");
            };
            Ok(handler.clone())
        };

        let acab = config.init_state()?;

        // Do everything else
        for (name, server) in &config.servers {
            match &server.variant {
                Variant::Prometheus(_) => {}
                Variant::Http(cfg) => {
                    let metrics = find_prometheus_server(&servers, name, &cfg.uses.metrics)?;
                    let handler = find_handler(&cfg.uses.handler_from)?;
                    let srv = freezer::Server {
                        bind: server.bind.clone(),
                        unix_socket_access: server.unix_socket_access,
                        variant: Bruce::new(
                            handler.language,
                            handler.compiler.as_ref(),
                            handler.path.as_ref(),
                            cfg.initial_seed.as_ref().unwrap_or(&config.initial_seed),
                            &metrics,
                            &acab.state.derive(name),
                            handler.config.clone(),
                        )?
                        .into(),
                    };
                    servers.insert(name.to_owned(), srv);
                }
                Variant::HaproxySPOA(cfg) => {
                    let metrics = find_prometheus_server(&servers, name, &cfg.uses.metrics)?;
                    let handler = find_handler(&cfg.uses.handler_from)?;
                    let srv = freezer::Server {
                        bind: server.bind.clone(),
                        unix_socket_access: server.unix_socket_access,
                        variant: CheckpointCharlie::new(
                            handler.language,
                            handler.compiler.as_ref(),
                            handler.path.as_ref(),
                            cfg.initial_seed.as_ref().unwrap_or(&config.initial_seed),
                            &metrics,
                            &acab.state.derive(name),
                            handler.config.clone(),
                        )?
                        .into(),
                    };
                    servers.insert(name.to_owned(), srv);
                }
            }
        }

        Ok(Self { servers })
    }
}