config.rs

use std::collections::BTreeMap;

use stackable_operator::{
    client::Client,
    k8s_openapi::api::core::v1::{
        ConfigMapKeySelector, EnvVar, EnvVarSource, SecretKeySelector, Volume, VolumeMount,
    },
    v2::types::kubernetes::NamespaceName,
};

use crate::{
    catalog::{FromTrinoCatalogError, ToCatalogConfig},
    crd::catalog::{TrinoCatalogConnector, TrinoCatalogName, v1alpha1},
};

#[derive(Clone, Debug)]
pub struct CatalogConfig {
    /// Name of the catalog
    pub name: TrinoCatalogName,

    /// Properties of the catalog
    pub properties: BTreeMap<String, String>,

    /// List of EnvVar that will be added to every Trino container
    pub env_bindings: Vec<EnvVar>,

    /// Env-Vars that should be exported.
    /// The value will be read from the file specified.
    /// You can think of it like `export <key>="$(cat <value>)"`
    pub load_env_from_files: BTreeMap<String, String>,

    /// Additional commands that needs to be executed before starting Trino
    pub init_container_extra_start_commands: Vec<String>,

    /// Volumes that need to be added to the pod (e.g. for S3 credentials)
    pub volumes: Vec<Volume>,

    /// Volume mounts that need to be added to the Trino container (e.g. for S3 credentials)
    pub volume_mounts: Vec<VolumeMount>,
}

impl CatalogConfig {
    pub fn new(name: &TrinoCatalogName, connector_name: impl Into<String>) -> Self {
        let mut config = CatalogConfig {
            name: name.clone(),
            properties: BTreeMap::new(),
            env_bindings: Vec::new(),
            load_env_from_files: BTreeMap::new(),
            init_container_extra_start_commands: Vec::new(),
            volumes: Vec::new(),
            volume_mounts: Vec::new(),
        };
        config.add_property("connector.name", connector_name);
        config
    }

    pub fn add_property(&mut self, property: impl Into<String>, value: impl Into<String>) {
        self.properties.insert(property.into(), value.into());
    }

    pub fn add_env_property_from_file(
        &mut self,
        property: impl Into<String>,
        file_name: impl Into<String>,
    ) {
        let property = property.into();
        let env_name = calculate_env_name(&self.name, &property);
        self.add_property(&property, format!("${{ENV:{env_name}}}"));
        self.load_env_from_files.insert(env_name, file_name.into());
    }

    pub fn add_env_property_from_secret(
        &mut self,
        property: impl Into<String>,
        secret_key_selector: SecretKeySelector,
    ) {
        let property = property.into();
        let env_name = calculate_env_name(&self.name, &property);
        self.env_bindings.push(EnvVar {
            name: env_name.clone(),
            value_from: Some(EnvVarSource {
                secret_key_ref: Some(secret_key_selector),
                ..Default::default()
            }),
            ..Default::default()
        });
        self.add_property(&property, format!("${{ENV:{env_name}}}"));
    }

    pub fn add_env_property_from_config_map(
        &mut self,
        property: impl Into<String>,
        config_map_key_selector: ConfigMapKeySelector,
    ) {
        let property = property.into();
        let env_name = calculate_env_name(&self.name, &property);
        self.env_bindings.push(EnvVar {
            name: env_name.clone(),
            value_from: Some(EnvVarSource {
                config_map_key_ref: Some(config_map_key_selector),
                ..Default::default()
            }),
            ..Default::default()
        });
        self.add_property(&property, format!("${{ENV:{env_name}}}"));
    }

    pub async fn from_catalog(
        catalog_name: &TrinoCatalogName,
        catalog: &v1alpha1::TrinoCatalog,
        client: &Client,
        catalog_namespace: &NamespaceName,
        trino_version: u16,
    ) -> Result<CatalogConfig, FromTrinoCatalogError> {
        let to_catalog_config: &dyn ToCatalogConfig = match &catalog.spec.connector {
            TrinoCatalogConnector::BlackHole(black_hole_connector) => black_hole_connector,
            TrinoCatalogConnector::DeltaLake(delta_lake_connector) => delta_lake_connector,
            TrinoCatalogConnector::GoogleSheet(google_sheet_connector) => google_sheet_connector,
            TrinoCatalogConnector::Generic(generic_connector) => generic_connector,
            TrinoCatalogConnector::Hive(hive_connector) => hive_connector,
            TrinoCatalogConnector::Iceberg(iceberg_connector) => iceberg_connector,
            TrinoCatalogConnector::Postgresql(postgresql_connector) => postgresql_connector,
            TrinoCatalogConnector::Tpcds(tpcds_connector) => tpcds_connector,
            TrinoCatalogConnector::Tpch(tpch_connector) => tpch_connector,
        };
        let mut catalog_config = to_catalog_config
            .to_catalog_config(catalog_name, catalog_namespace, client, trino_version)
            .await?;

        catalog_config
            .properties
            .extend(catalog.spec.config_overrides.clone());

        for removal in &catalog.spec.config_removals {
            if catalog_config.properties.remove(removal).is_none() {
                tracing::warn!(
                    catalog.name = %catalog_name,
                    property = removal,
                    "You asked to remove a non-existing config property from a catalog"
                );
            }
        }

        Ok(catalog_config)
    }
}

fn calculate_env_name(catalog_name: &TrinoCatalogName, property: impl Into<String>) -> String {
    let catalog = catalog_name.to_string().replace(['.', '-'], "_");
    let property = property.into().replace(['.', '-'], "_");
    format!("CATALOG_{catalog}_{property}").to_uppercase()
}