use clap::Parser;

use tempfile::{Builder, TempDir};

use tracing_subscriber::prelude::*;
use tracing_tracy::TracyLayer;

use lsm::{Database, Params, StartMode, WriteOptions};

#[derive(Parser)]
#[clap(author, version, about, long_about = None)]
struct Args {
    #[clap(long)]
    enable_tracing: bool,
    #[clap(long)]
    log_level_stats: Option<String>,
    #[clap(long, default_value = "100000")]
    num_entries: usize,
}

async fn bench_init(args: &Args) -> (TempDir, Database) {
    if args.enable_tracing {
        tracing_subscriber::registry()
            .with(TracyLayer::default())
            .init();
    }

    let _ = env_logger::builder().is_test(true).try_init();
    let tmp_dir = Builder::new()
        .prefix("lsm-async-benchmark-")
        .tempdir()
        .unwrap();

    let mut db_path = tmp_dir.path().to_path_buf();
    db_path.push("storage.lsm");

    let params = Params {
        db_path,
        log_level_stats: args.log_level_stats.clone(),
        ..Default::default()
    };
    const SM: StartMode = StartMode::CreateOrOverride;

    let database = Database::new_with_params(SM, params)
        .await
        .expect("Failed to create database instance");

    (tmp_dir, database)
}

async fn main_inner() {
    let args = Args::parse();

    let (_tmpdir, database) = bench_init(&args).await;

    log::info!("Starting read/write benchmark");

    let options = WriteOptions { sync: false };

    log::debug!("Writing {} entries", args.num_entries);

    for pos in 0..args.num_entries {
        let key = format!("{pos}").into_bytes();
        let value = format!("some_string_{pos}").into_bytes();
        database.put_opts(key, value, &options).await.unwrap();
    }

    log::debug!("Reading {} entries", args.num_entries);

    for pos in 0..args.num_entries {
        let key = format!("{pos}").into_bytes();
        let expected = format!("some_string_{pos}").into_bytes();

        assert_eq!(
            database.get(&key).await.unwrap().unwrap().get_value(),
            expected,
        );
    }

    database.stop().await.unwrap();
    log::info!("Done");
}

#[cfg(feature = "async-io")]
fn main() {
    tokio_uring::start(async {
        main_inner().await;
    });
}

#[cfg(not(feature = "async-io"))]
#[tokio::main]
async fn main() {
    main_inner().await;
}