cuddle-clusters/crates/cuddle-clusters/src/process.rs

use std::{
    collections::HashMap,
    ops::Deref,
    path::{Path, PathBuf},
};

use anyhow::Context;
use futures::{future::BoxFuture, FutureExt};
use minijinja::context;
use tokio::io::AsyncWriteExt;
use tokio_stream::{wrappers::ReadDirStream, StreamExt};

use crate::components::{ConcreteComponent, IntoComponent};

pub async fn process() -> anyhow::Result<()> {
    process_opts(
        Vec::<ConcreteComponent>::new(),
        ProcessOpts::default(),
        None::<NoUploadStrategy>,
    )
    .await
}

pub trait UploadStrategy {
    fn upload(&self, input_path: &Path) -> BoxFuture<'_, anyhow::Result<()>>;
}

#[derive(Default)]
pub struct NoUploadStrategy {}

impl UploadStrategy for NoUploadStrategy {
    fn upload(&self, _input_path: &Path) -> BoxFuture<'_, anyhow::Result<()>> {
        async move { Ok(()) }.boxed()
    }
}

pub struct ProcessOpts {
    pub path: PathBuf,
    pub output: PathBuf,

    pub variables: HashMap<String, String>,
}

impl Default for ProcessOpts {
    fn default() -> Self {
        Self {
            path: std::env::current_dir().expect("to be able to get current dir"),
            output: std::env::current_dir()
                .expect("to be able to get current dir")
                .join("cuddle-clusters")
                .join("k8s"),
            variables: HashMap::default(),
        }
    }
}

const TEMPLATES_PATH_PREFIX: &str = "templates/clusters";
const CUDDLE_PLAN_PATH_PREFIX: &str = ".cuddle/base";

pub async fn process_opts(
    components: impl IntoIterator<Item = impl IntoComponent>,
    opts: ProcessOpts,
    upload_strategy: Option<impl UploadStrategy>,
) -> anyhow::Result<()> {
    let components = components
        .into_iter()
        .map(|c| c.into_component())
        .collect::<Vec<_>>();

    let path = opts.path.canonicalize().context("failed to find folder")?;

    let cuddle_path = path.join("cuddle.yaml");

    tracing::debug!(
        "searching for templates in: {} with cuddle: {}",
        path.display(),
        cuddle_path.display()
    );

    let clusters = read_cuddle_section(&cuddle_path).await?;
    tracing::debug!("found clusters: {:?}", clusters);

    let template_files = load_template_files(&path).await?;
    tracing::debug!("found files: {:?}", template_files);

    let _ = tokio::fs::remove_dir_all(&opts.output).await;
    tokio::fs::create_dir_all(&opts.output).await?;

    process_templates(
        &components,
        &clusters,
        &template_files,
        &opts.output,
        &opts.variables,
    )
    .await?;

    if let Some(upload_strategy) = upload_strategy {
        upload_strategy.upload(&opts.output).await?;
    }

    Ok(())
}

#[derive(serde::Deserialize, Default, Debug, Clone)]
struct CuddleClusters(HashMap<String, serde_yaml::Value>);

impl Deref for CuddleClusters {
    type Target = HashMap<String, serde_yaml::Value>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

async fn read_cuddle_section(path: &Path) -> anyhow::Result<CuddleClusters> {
    let cuddle_file = tokio::fs::read(path)
        .await
        .context(format!("failed to read: {}", path.display()))?;

    let value: serde_yaml::Value = serde_yaml::from_slice(&cuddle_file)?;

    let mut cuddle_clusters = CuddleClusters::default();
    if let Some(clusters) = value.get("cuddle/clusters") {
        if let Some(mapping) = clusters.as_mapping() {
            for (key, value) in mapping.iter() {
                if let Some(key) = key.as_str() {
                    cuddle_clusters.0.insert(key.to_string(), value.clone());
                }
            }
        }
    }

    Ok(cuddle_clusters)
}

#[derive(Debug, Clone, Default)]
struct TemplateFiles {
    templates: HashMap<String, PathBuf>,
    raw: HashMap<String, PathBuf>,
}

impl TemplateFiles {
    fn merge(&mut self, template_files: TemplateFiles) {
        for (name, path) in template_files.templates {
            // Ignored if the key is already present
            let _ = self.templates.try_insert(name, path);
        }
        for (name, path) in template_files.raw {
            // Ignored if the key is already present
            let _ = self.raw.try_insert(name, path);
        }
    }
}

fn load_template_files(path: &Path) -> BoxFuture<'static, anyhow::Result<TemplateFiles>> {
    let template_path = path.join(TEMPLATES_PATH_PREFIX);
    let path = path.to_path_buf();

    async move {
        if template_path.exists() {
            let templates = read_dir(&template_path)
                .await?
                .into_iter()
                .filter(|(_, i)| i.extension().and_then(|e| e.to_str()) == Some("jinja2"))
                .collect();
            let raw = read_dir(&template_path.join("raw"))
                .await
                .unwrap_or_default();
            let mut template_files = TemplateFiles { templates, raw };

            let nested_path = path.join(CUDDLE_PLAN_PATH_PREFIX);
            if nested_path.exists() {
                let nested = load_template_files(&nested_path).await?;
                template_files.merge(nested);
            }
            Ok(template_files)
        } else {
            let nested_path = path.join(CUDDLE_PLAN_PATH_PREFIX);
            if nested_path.exists() {
                let nested = load_template_files(&nested_path).await?;
                Ok(nested)
            } else {
                Ok(TemplateFiles::default())
            }
        }
    }
    .boxed()
}

async fn read_dir(path: &Path) -> anyhow::Result<HashMap<String, PathBuf>> {
    let template_dir = tokio::fs::read_dir(path).await?;
    let mut template_dir_stream = ReadDirStream::new(template_dir);

    let mut paths = HashMap::new();
    while let Some(entry) = template_dir_stream.next().await {
        let entry = entry?;

        if entry.metadata().await?.is_file() {
            paths.insert(
                entry
                    .path()
                    .file_name()
                    .expect("the file to have a filename")
                    .to_string_lossy()
                    .to_string(),
                entry.path(),
            );
        }
    }

    Ok(paths)
}

async fn process_templates(
    components: &Vec<ConcreteComponent>,
    clusters: &CuddleClusters,
    template_files: &TemplateFiles,
    dest: &Path,
    variables: &HashMap<String, String>,
) -> anyhow::Result<()> {
    for (environment, value) in clusters.iter() {
        process_cluster(
            components,
            value,
            environment,
            template_files,
            &dest.join(environment),
            variables,
        )
        .await?;
    }

    Ok(())
}

async fn process_cluster(
    components: &Vec<ConcreteComponent>,
    value: &serde_yaml::Value,
    environment: &str,
    template_files: &TemplateFiles,
    dest: &Path,
    variables: &HashMap<String, String>,
) -> anyhow::Result<()> {
    for (_, template_file) in &template_files.templates {
        process_template_file(
            components,
            value,
            environment,
            template_file,
            dest,
            variables,
        )
        .await?;
    }

    for (template_file_name, template_content) in components
        .iter()
        .filter_map(|c| c.render(environment, value))
        .flat_map(|v| match v {
            Ok(v) => Ok(v),
            Err(e) => {
                tracing::warn!("failed to render value for template");
                Err(e)
            }
        })
    {
        process_render_template(
            components,
            value,
            environment,
            &template_file_name,
            &template_content,
            dest,
            variables,
        )
        .await?;
    }

    for (_, raw_file) in &template_files.raw {
        process_raw_file(environment, raw_file, dest).await?;
    }

    Ok(())
}

async fn process_template_file(
    components: &Vec<ConcreteComponent>,
    value: &serde_yaml::Value,
    environment: &str,
    template_file: &PathBuf,
    dest: &Path,
    variables: &HashMap<String, String>,
) -> anyhow::Result<()> {
    let file = tokio::fs::read_to_string(template_file)
        .await
        .context(format!("failed to find file: {}", template_file.display()))?;
    let file_name = template_file
        .file_stem()
        .ok_or(anyhow::anyhow!("file didn't have a jinja2 format"))?;

    process_render_template(
        components,
        value,
        environment,
        &file_name.to_string_lossy(),
        &file,
        dest,
        variables,
    )
    .await?;

    Ok(())
}

async fn process_render_template(
    components: &Vec<ConcreteComponent>,
    value: &serde_yaml::Value,
    environment: &str,
    file_name: &str,
    file_content: &str,
    dest: &Path,
    user_vars: &HashMap<String, String>,
) -> anyhow::Result<()> {
    if !dest.exists() {
        tokio::fs::create_dir_all(dest).await?;
    }

    let mut dest_file = tokio::fs::File::create(dest.join(file_name)).await?;

    let mut env = minijinja::Environment::new();
    env.add_template(file_name, &file_content)
        .context(format!("failed to load template at: {}", file_name))?;
    env.add_global("environment", environment);

    let mut variables = HashMap::new();

    for component in components {
        let name = component.name();

        if let Some(value) = component.render_value(environment, value) {
            let value = value?;

            variables.insert(name.replace("/", "_"), value);
        }
    }
    variables.insert(
        "user_vars".into(),
        minijinja::Value::from_serialize(user_vars),
    );

    let tmpl = env.get_template(file_name)?;
    let rendered = tmpl.render(context! {
        vars => variables
    })?;

    let rendered = if rendered.is_empty() || rendered.ends_with("\n") {
        rendered
    } else {
        format!("{rendered}\n")
    };

    dest_file.write_all(rendered.as_bytes()).await?;

    Ok(())
}

async fn process_raw_file(
    _cluster_name: &str,
    raw_file: &PathBuf,
    dest: &Path,
) -> anyhow::Result<()> {
    let file = tokio::fs::read_to_string(raw_file)
        .await
        .context(format!("failed to find file: {}", raw_file.display()))?;

    if !dest.exists() {
        tokio::fs::create_dir_all(dest).await?;
    }

    let file_name = raw_file
        .file_name()
        .ok_or(anyhow::anyhow!("file didn't have a file name"))?;

    let mut dest_file = tokio::fs::File::create(dest.join(file_name)).await?;
    dest_file.write_all(file.as_bytes()).await?;

    Ok(())
}