BoosAPI/src/main.rs

use anyhow::Context;
use axum::serve;
use boosapi::{
    config::AppConfig,
    db,
    services::{builtin_keys, runtime_auth},
    AppState,
};
use clap::{Parser, Subcommand};
use std::ffi::OsString;
use std::net::SocketAddr;
use tokio::signal;
use tracing_subscriber::{prelude::*, EnvFilter};

#[derive(Parser, Debug)]
#[command(
    name = "boosapi",
    version,
    author,
    about = "AI API Gateway — Rust rewrite of BoosAPI"
)]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,

    /// Config file path
    #[arg(short, long, default_value = "config.toml")]
    config: String,

    /// Host to bind to
    #[arg(long)]
    host: Option<String>,

    /// Port to listen on
    #[arg(short, long)]
    port: Option<u16>,

    /// Log level (trace, debug, info, warn, error)
    #[arg(long)]
    log_level: Option<String>,

    /// Only run database migrations and exit
    #[arg(long)]
    migrate_only: bool,

    /// Load and validate configuration, then exit
    #[arg(long)]
    check_config: bool,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Copy core data from SQLite into PostgreSQL and exit
    MigrateSqliteToPostgres {
        /// Source SQLite URL, e.g. sqlite:boosapi.db?mode=rwc
        #[arg(long)]
        sqlite_url: String,

        /// Target PostgreSQL URL, e.g. postgres://user:pass@host:5432/boosapi
        #[arg(long)]
        postgres_url: String,

        /// Truncate target business tables before importing
        #[arg(long)]
        force_truncate: bool,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Cli::parse_from(normalize_legacy_args());

    if let Some(Commands::MigrateSqliteToPostgres {
        sqlite_url,
        postgres_url,
        force_truncate,
    }) = &args.command
    {
        let sqlite = db::create_pool(sqlite_url, 1).await?;
        let pg = db::create_pg_pool(postgres_url, 2).await?;
        db::run_migrations(db::DbKind::Postgres, None, Some(&pg)).await?;
        let summary = db::sqlite_to_postgres::migrate(&sqlite, &pg, *force_truncate).await?;
        println!("{}", serde_json::to_string_pretty(&summary)?);
        return Ok(());
    }

    // Load config
    let mut config = AppConfig::load_from_file(&args.config)?;

    // Apply CLI overrides
    if let Some(host) = args.host {
        config.server.host = host;
    }
    if let Some(port) = args.port {
        config.server.port = port;
    }
    if let Some(level) = args.log_level {
        config.server.log_level = level;
    }

    if args.check_config {
        println!("Configuration is valid.");
        println!("Server: {}:{}", config.server.host, config.server.port);
        println!("Database URL: {}", config.db.url);
        return Ok(());
    }

    // Initialize tracing
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new(&config.server.log_level));

    tracing_subscriber::registry()
        .with(filter)
        .with(tracing_subscriber::fmt::layer())
        .with(boosapi::services::runtime_logs::RuntimeLogLayer)
        .init();

    tracing::info!("Loaded configuration from {}", args.config);

    let db_kind = db::DbKind::from_url(&config.db.url)?;

    // Create database pool
    let pool = db::create_pool(&config.db.url, config.db.max_connections).await?;
    tracing::info!("Database connected");

    // Run migrations
    db::run_migrations(db_kind, Some(&pool), None).await?;
    db::schema_compat::ensure_sqlite_upstream_compat(&pool).await?;
    tracing::info!("Database migrations applied");

    if args.migrate_only {
        tracing::info!("Migration only mode requested. Exiting.");
        return Ok(());
    }

    // Bootstrap built-in system keys (idempotent).
    let comfyui_agent_key = builtin_keys::ensure_comfyui_agent_key(&pool)
        .await
        .context("failed to bootstrap built-in ComfyUI Agent key")?;
    tracing::info!(
        key_id = comfyui_agent_key.0,
        "Built-in ComfyUI Agent downstream key ready"
    );
    let builtin = builtin_keys::BuiltinKeys {
        comfyui_agent: comfyui_agent_key,
    };
    let runtime_auth =
        runtime_auth::RuntimeAuthSettings::load(&pool, config.auth.proxy_token.clone())
            .await
            .context("failed to load runtime auth settings")?;

    // Build app state
    let state = AppState::new(config.clone(), pool, builtin, runtime_auth);
    if let Err(err) = state.comfyui.start().await {
        tracing::warn!(%err, "ComfyUI embedded process did not start");
    }
    let comfyui_process = state.comfyui.clone();

    // Build router
    let app = boosapi::routes::build_router(state);

    // Start server
    let addr = SocketAddr::new(config.server.host.parse()?, config.server.port);
    tracing::info!("Starting server on {addr}");

    serve(
        tokio::net::TcpListener::bind(addr).await?,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown_signal())
    .await?;

    if let Err(err) = comfyui_process.stop().await {
        tracing::warn!(%err, "ComfyUI embedded process did not stop cleanly");
    }

    Ok(())
}

fn normalize_legacy_args() -> Vec<OsString> {
    std::env::args_os()
        .map(|arg| {
            if arg == "-help" {
                OsString::from("--help")
            } else if arg == "-version" {
                OsString::from("--version")
            } else {
                arg
            }
        })
        .collect()
}

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    tracing::info!("Shutdown signal received, starting graceful shutdown");
}