YANPM/apps/container/src/main.rs

use std::sync::Arc;

use clap::Parser;
use container::{
    Config,
    containers::{
        ConfigInfoType,
        agent::{AgentConfig, AgentContainerConfig, AgentContainerInfo, NginxConfig},
        db::DBInfo,
    },
    start_attached,
};

/// Command line arguments
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Database type to use: 'postgres' or 'sqlite'. Can also be set with DB_TYPE env var.
    #[arg(long, default_value = "sqlite", env = "DB_TYPE")]
    db_type: String,

    // agent related
    /// agent image name
    #[arg(long, default_value = "yanpm/agent", env = "AGENT_IMAGE_NAME")]
    agent_image: String,
    /// agent image tag
    #[arg(long, default_value = "latest", env = "AGENT_IMAGE_TAG")]
    agent_image_tag: String,
    /// force build agent image
    #[arg(long, default_value_t = false, env = "AGENT_FORCE_BUILD")]
    agent_force_build: bool,
    /// dockerfile path for building agent image
    #[arg(long, env = "AGENT_DOCKERFILE_PATH", required = false)]
    agent_dockerfile_path: Option<String>,
    /// host's location to mount nginx config files folder generated by the agent
    #[arg(long, env = "AGENT_NGINX_CONFIG_DIR", required = false)]
    agent_nginx_config_dir: Option<String>,
    /// host's location folder to mount the unix socket files
    #[arg(long, env = "AGENT_SOCK_PATH", required = false)]
    agent_sock_path: Option<String>,
    /// socket permissions to set on the unix socket
    #[arg(long, default_value = "660", env = "AGENT_SOCK_PERM", required = false)]
    agent_sock_perm: u32,
    /// socket GID to set on the unix socket
    #[arg(long, default_value = "", env = "AGENT_SOCK_GID", required = false)]
    agent_sock_gid: String,
    /// nginx expose http port
    #[arg(
        long,
        default_value_t = true,
        env = "AGENT_NGINX_EXPOSE_HTTP",
        required = false
    )]
    agent_nginx_expose_http: bool,
    /// nginx expose https port
    #[arg(
        long,
        default_value_t = false,
        env = "AGENT_NGINX_EXPOSE_HTTPS",
        required = false
    )]
    agent_nginx_expose_https: bool,
}

struct ParsedArgs {
    db_type: String,
    agent_container_config: Option<AgentContainerConfig>,
}

#[tokio::main]
async fn main() {
    let args = parse_args().await;

    println!("Starting container with database type: {}", args.db_type);
    let db_config = match args.db_type.to_lowercase().as_str() {
        "postgres" | "pg" | "pgsql" => {
            use container::containers::db::postgresql::PostgreSQLContainer;
            println!("Using PostgreSQL database");
            PostgreSQLContainer::new(None)
                .await
                .get_db_container_config_info()
                .await
        }
        "sqlite" | "sql" => {
            println!("Using SQLite database");
            use container::containers::db::sqlite::SQLiteContainer;
            SQLiteContainer::new(None)
                .await
                .get_db_container_config_info()
                .await
        }
        other => {
            eprintln!("Unknown db_type: {}. Use 'postgres' or 'sqlite'", other);
            std::process::exit(1);
        }
    };
    println!("Database configuration obtained.");

    let agent_container = if let Some(agent_config) = &args.agent_container_config {
        println!(
            "Agent container will be used with socket folder: {} and nginx config dir: {}",
            agent_config.agent_config.sock_folder, agent_config.agent_config.nginx_config_dir
        );
        Some(agent_config.get_unstarted_container().await)
    } else {
        println!("No agent container configuration provided, skipping agent setup.");
        None
    };

    let config = Config {
        database: db_config,
        agent: match agent_container {
            Some(Ok(container)) => Some(ConfigInfoType::Containerized(AgentContainerInfo {
                container: Arc::new(container.await.expect("Failed to start agent container")),
                config: args.agent_container_config.expect("Invalid config state"),
            })),
            Some(Err(e)) => {
                eprintln!("Failed to set up agent container: {}", e);
                std::process::exit(1);
            }
            None => None,
        },
    };

    println!("Starting container...");
    start_attached(&config).await;
    println!("Container stopped. Exiting...");
}

async fn parse_args() -> ParsedArgs {
    // Parse command line arguments and environment variables
    let args = Args::parse();

    // if any required args are missing, do not start agent
    let dockerfile_path = match args.agent_dockerfile_path {
        None => {
            println!("Agent dockerfile path not provided, skipping agent setup.");
            return ParsedArgs {
                db_type: args.db_type,
                agent_container_config: None,
            };
        }
        Some(path) => path,
    };

    let time = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs();

    let agent_config = AgentConfig {
        sock_folder: match args.agent_sock_path {
            None => {
                // create a temp dir for the socket path
                let temp_dir = std::env::temp_dir().join(format!("yanpm-agent-sock-{}", time));
                std::fs::create_dir_all(&temp_dir)
                    .expect("Failed to create temp dir for agent socket");
                temp_dir.to_string_lossy().to_string()
            }
            Some(path) => path,
        },
        nginx_config_dir: match args.agent_nginx_config_dir {
            None => {
                // create a temp dir for the nginx config dir
                let temp_dir =
                    std::env::temp_dir().join(format!("yanpm-agent-nginx-configs-{}", time));
                std::fs::create_dir_all(&temp_dir)
                    .expect("Failed to create temp dir for agent nginx configs");
                temp_dir.to_string_lossy().to_string()
            }
            Some(path) => path,
        },
        sock_perm: args.agent_sock_perm,
        sock_gid: args.agent_sock_gid.clone(),
    };

    ParsedArgs {
        db_type: args.db_type,
        agent_container_config: Some(AgentContainerConfig {
            // TODO: allow customization of these fields via CLI args
            image: args.agent_image,
            tag: args.agent_image_tag,
            container_name: format!("yanpm-agent-container-{}", time),
            dockerfile_path,
            force_build: args.agent_force_build,
            agent_config,
            nginx_config: NginxConfig {
                expose_http: args.agent_nginx_expose_http,
                expose_https: args.agent_nginx_expose_https,
            },
        }),
    }
}