use clap::{Parser, Subcommand};

#[derive(Parser)]
#[command(name = "mkv", about = "Distributed key-value store")]
struct Cli {
    #[arg(short, long, default_value = "/tmp/mkv/index.db")]
    db: String,

    #[arg(short, long, required = true, value_delimiter = ',')]
    volumes: Vec<String>,

    #[arg(short, long, default_value_t = 2)]
    replicas: usize,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Start the index server
    Serve {
        #[arg(short, long, default_value_t = 3000)]
        port: u16,
    },
    /// Rebuild SQLite index from volume servers
    Rebuild,
    /// Rebalance data after adding/removing volumes
    Rebalance {
        #[arg(long)]
        dry_run: bool,
    },
}

async fn shutdown_signal() {
    let ctrl_c = tokio::signal::ctrl_c();
    #[cfg(unix)]
    {
        let mut sigterm =
            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                .expect("failed to install SIGTERM handler");
        tokio::select! {
            _ = ctrl_c => tracing::info!("Received SIGINT, shutting down..."),
            _ = sigterm.recv() => tracing::info!("Received SIGTERM, shutting down..."),
        }
    }
    #[cfg(not(unix))]
    {
        ctrl_c.await.expect("failed to listen for ctrl-c");
        tracing::info!("Received ctrl-c, shutting down...");
    }
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    let cli = Cli::parse();

    let args = mkv::Args {
        db_path: cli.db,
        volumes: cli.volumes,
        replicas: cli.replicas,
    };

    match cli.command {
        Commands::Serve { port } => {
            let app = mkv::build_app(&args);
            let addr = format!("0.0.0.0:{port}");
            tracing::info!("Listening on {addr}");
            let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
            axum::serve(listener, app)
                .with_graceful_shutdown(shutdown_signal())
                .await
                .unwrap();
            tracing::info!("Shutdown complete");
        }
        Commands::Rebuild => {
            mkv::rebuild::run(&args).await;
        }
        Commands::Rebalance { dry_run } => {
            mkv::rebalance::run(&args, dry_run).await;
        }
    }
}