use clap::Parser;
use rocksdb::{CompactOptions, DB, DBCompressionType, Options};

use std::path::{Path, PathBuf};
use std::sync::{Arc, atomic::Ordering};

use space_efficiency_check::work::{ProcessError, run_workers};

/// import a directory of car files into rocks to check space efficiency
#[derive(Parser)]
struct Args {
    /// directory of car files
    car_dir: PathBuf,

    /// where to kick rocks
    #[arg(long, default_value = "./local/space-efficiency-check-db")]
    db_dir: PathBuf,

    /// concurrent car workers
    #[arg(long, default_value_t = 10)]
    workers: usize,

    /// repo-stream memory limit: we don't disk spill; set higher than largest car
    #[arg(long, default_value_t = 1024)]
    mem_limit_mb: usize,

    /// skip import and just do manual compact
    #[arg(long, action)]
    just_compact: bool,

    /// sample fraction of DIDs to import (0.0–1.0, default: all)
    #[arg(long)]
    sample: Option<f64>,

    /// store sequential u64s in place of dids
    #[arg(long, action)]
    intern: bool,
}

fn open_db(path: &Path) -> Result<DB, rocksdb::Error> {
    let mut opts = Options::default();
    opts.create_if_missing(true);

    opts.set_compression_type(DBCompressionType::Lz4);

    opts.set_max_subcompactions(4);

    let mut block_opts = rocksdb::BlockBasedOptions::default();
    block_opts.set_block_size(32_768);
    opts.set_block_based_table_factory(&block_opts);

    DB::open(&opts, path)
}

#[tokio::main]
async fn main() -> Result<(), ProcessError> {
    tracing_subscriber::fmt::init();

    let args = Args::parse();

    let db = Arc::new(open_db(&args.db_dir)?);

    if args.just_compact {
        manually_compact(db)?;
        return Ok(());
    }

    let stats = run_workers(
        &args.car_dir,
        db.clone(),
        args.workers,
        args.mem_limit_mb,
        args.sample,
        args.intern,
    )
    .await?;

    let repos = stats.repos.load(Ordering::Relaxed);
    let empty = stats.empty_repos.load(Ordering::Relaxed);
    let records = stats.records.load(Ordering::Relaxed);
    let failed = stats.failed_repos.load(Ordering::Relaxed);
    let car_bytes = stats.car_bytes.load(Ordering::Relaxed);

    tracing::info!(repos, empty, records, failed, "import complete");
    tracing::info!(
        car_bytes,
        car_mib = car_bytes / 2_u64.pow(20),
        "total input size"
    );

    let size_before = db_size(&db);
    tracing::info!(
        bytes = size_before,
        mib = size_before / 2_u64.pow(20),
        "db size before compaction"
    );

    let size_after = manually_compact(db)?;

    if car_bytes > 0 {
        let ratio = size_after as f64 / car_bytes as f64;
        tracing::info!(ratio = format!("{ratio:.3}"), "db size / car bytes");
    }

    Ok(())
}

fn manually_compact(db: Arc<DB>) -> Result<u64, ProcessError> {
    let mut compact_opts = CompactOptions::default();
    compact_opts.set_bottommost_level_compaction(rocksdb::BottommostLevelCompaction::Force);

    tracing::info!("running manual compaction...");
    let compact_start = std::time::Instant::now();
    db.compact_range_opt(None::<&[u8]>, None::<&[u8]>, &compact_opts);
    let compact_elapsed = compact_start.elapsed();

    let size_after = db_size(&db);
    tracing::info!(
        bytes = size_after,
        mib = size_after / 2_u64.pow(20),
        ?compact_elapsed,
        "db size after compaction",
    );
    Ok(size_after)
}

fn db_size(db: &DB) -> u64 {
    db.property_value("rocksdb.total-sst-files-size")
        .ok()
        .flatten()
        .and_then(|v| v.parse::<u64>().ok())
        .unwrap_or(0)
}