mod config;
mod db;
mod error;
mod hasher;
mod health;
mod server;
mod volume;

use clap::{Parser, Subcommand};
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Parser)]
#[command(name = "mkv", about = "Distributed key-value store")]
struct Cli {
    #[arg(short, long, default_value = "config.toml")]
    config: PathBuf,

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

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

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();

    let cli = Cli::parse();
    let config = config::Config::load(&cli.config).unwrap_or_else(|e| {
        eprintln!("Failed to load config: {e}");
        std::process::exit(1);
    });

    match cli.command {
        Commands::Serve => serve(config).await,
        Commands::Rebuild => {
            eprintln!("rebuild not yet implemented");
            std::process::exit(1);
        }
        Commands::Rebalance { dry_run: _ } => {
            eprintln!("rebalance not yet implemented");
            std::process::exit(1);
        }
    }
}

async fn serve(config: config::Config) {
    let db_path = &config.database.path;

    // Ensure parent directory exists
    if let Some(parent) = std::path::Path::new(db_path).parent() {
        std::fs::create_dir_all(parent).unwrap_or_else(|e| {
            eprintln!("Failed to create database directory: {e}");
            std::process::exit(1);
        });
    }

    let (writer, ready_rx) = db::spawn_writer(db_path.to_string());
    ready_rx.await.expect("writer failed to initialize");

    let num_readers = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4);
    let reads = db::ReadPool::new(db_path, num_readers);

    let volume_urls = config.volume_urls();
    let ring = Arc::new(RwLock::new(hasher::Ring::new(
        &volume_urls,
        config.server.virtual_nodes,
    )));

    let volume_client = volume::VolumeClient::new();

    // Start with all volumes assumed healthy, health checker will update
    let healthy_volumes: health::HealthyVolumes =
        Arc::new(RwLock::new(HashSet::from_iter(volume_urls.clone())));

    health::spawn_health_checker(
        volume_client.clone(),
        volume_urls.clone(),
        healthy_volumes.clone(),
    );

    let port = config.server.port;

    let state = server::AppState {
        writer,
        reads,
        ring,
        volume_client,
        healthy_volumes,
        config: Arc::new(config),
    };

    tracing::info!("Starting mkv server on port {port}");
    tracing::info!("  Readers: {num_readers}");
    tracing::info!("  Volumes: {volume_urls:?}");
    tracing::info!("  Replication factor: {}", state.config.server.replication_factor);

    let app = axum::Router::new()
        .route("/", axum::routing::get(server::list_keys))
        .route(
            "/{*key}",
            axum::routing::get(server::get_key)
                .put(server::put_key)
                .delete(server::delete_key)
                .head(server::head_key),
        )
        .with_state(state);

    let addr = format!("0.0.0.0:{port}");
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    tracing::info!("Listening on {addr}");
    axum::serve(listener, app).await.unwrap();
}