orign 0.2.3

A globally distributed container orchestrator
Documentation
mod cli;
mod commands;
mod models;

use crate::cli::{
    Cli, Commands, CreateCommands, DeleteCommands, DeploymentCommands, GetCommands, StopCommands,
    TrainCommands, WorkCommands,
};
use clap::Parser;
use orign::db;
use std::error::Error;
use tracing_subscriber;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Initialize tracing
    tracing_subscriber::fmt::init();

    // Parse command-line arguments
    let cli = Cli::parse();

    match cli.command {
        Commands::Serve { host, port } => {
            commands::serve_cmd::execute(host, port).await?;
        }
        Commands::Chat {
            model,
            msg,
            image,
            framework,
            adapter,
        } => {
            commands::chat_cmd::execute(model, msg, image, adapter, framework).await?;
        }
        Commands::Trigger { buffer_name } => {
            commands::trigger_cmd::trigger_buffer(buffer_name).await?;
        }
        Commands::Create { command } => match command {
            CreateCommands::Buffer { command } => {
                commands::create_cmd::create_replay_buffer(command).await?;
            }
        },
        Commands::Get { command } => match command {
            GetCommands::Deployments { id } => {
                commands::get_cmd::get_deployment(id).await?;
            }
            GetCommands::Trainings { id } => {
                commands::get_cmd::get_training(id).await?;
            }
            GetCommands::Buffers { name } => {
                commands::get_cmd::get_buffer(name).await?;
            }
            GetCommands::Models {} => {
                commands::get_cmd::get_models().await?;
            }
            GetCommands::Datasets {} => {
                commands::get_cmd::get_datasets().await?;
            }
            GetCommands::Adapters {} => {
                commands::get_cmd::get_adapters().await?;
            }
        },
        Commands::Delete { command } => match command {
            DeleteCommands::Deployment { id } => {
                commands::delete_cmd::delete_deployment(id).await?;
            }
            DeleteCommands::Buffer { name } => {
                commands::delete_cmd::delete_buffer(name).await?;
            }
            DeleteCommands::Adapter { name } => {
                commands::delete_cmd::delete_adapter(name).await?;
            }
        },
        Commands::Stop { command } => match command {
            StopCommands::Training { id } => {
                commands::train_cmd::stop_training(id).await?;
            }
        },
        Commands::Send {
            buffer_name,
            file,
            train,
        } => {
            commands::send_cmd::send_buffer(buffer_name, file, train).await?;
        }
        Commands::Prepare {
            dataset_type,
            url,
            split_ratio,
            base_path,
            image_path,
            num_workers,
        } => {
            commands::prepare_cmd::execute_prepare(
                dataset_type,
                url,
                split_ratio,
                base_path,
                image_path,
                num_workers,
            )
            .await?;
        }
        Commands::Train { command } => match command {
            TrainCommands::Swift(args) => {
                commands::train_cmd::run_swift_training(args).await?;
            }
        },
        Commands::Report {
            watch_dir,
            save_dir,
            debounce_secs,
        } => {
            commands::report_cmd::execute_report(watch_dir, save_dir, debounce_secs).await?;
        }
        Commands::Work { command } => match command {
            WorkCommands::ReportTrainings { k8s_namespace } => {
                // 1) Obtain or build your DB connection / AppState
                //    You might have a helper function to build the db, or retrieve AppState, etc.
                //    e.g.:
                // let app_state = commands::serve_cmd::build_app_state().await?;
                // Or maybe you have a direct DB connection getter:
                // let db = commands::some_cmd::get_db_connection().await?;
                let db = db::init_db().await?;
                println!("Database connected successfully.");

                // 2) Call your logic to report training statuses from K8s
                commands::work_cmd::execute_report_jobs(&db, &k8s_namespace).await?;
            }
        },
        Commands::Login => {
            commands::login_cmd::execute().await?;
        }
    }

    Ok(())
}