use rocksdb::{DB, IteratorMode, Options};
use rpki::storage::{column_family_descriptors, CF_REPOSITORY_VIEW};
use std::fs;
use std::path::{Path, PathBuf};

fn usage() -> String {
    let bin = "repository_view_stats";
    format!(
        "\
Usage:
  {bin} --db <path>

Options:
  --db <path>   RocksDB directory
  --help        Show this help
"
    )
}

fn parse_args(argv: &[String]) -> Result<PathBuf, String> {
    if argv.iter().any(|arg| arg == "--help" || arg == "-h") {
        return Err(usage());
    }
    let mut db_path: Option<PathBuf> = None;
    let mut i = 1usize;
    while i < argv.len() {
        match argv[i].as_str() {
            "--db" => {
                i += 1;
                let value = argv.get(i).ok_or("--db requires a value")?;
                db_path = Some(PathBuf::from(value));
            }
            other => return Err(format!("unknown argument: {other}\n\n{}", usage())),
        }
        i += 1;
    }
    db_path.ok_or_else(|| format!("--db is required\n\n{}", usage()))
}

fn dir_size(path: &Path) -> Result<u64, Box<dyn std::error::Error>> {
    let mut total = 0u64;
    for entry in fs::read_dir(path)? {
        let entry = entry?;
        let file_type = entry.file_type()?;
        if file_type.is_file() {
            total = total.saturating_add(entry.metadata()?.len());
        } else if file_type.is_dir() {
            total = total.saturating_add(dir_size(&entry.path())?);
        }
    }
    Ok(total)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let argv: Vec<String> = std::env::args().collect();
    let db_path = parse_args(&argv).map_err(|e| -> Box<dyn std::error::Error> { e.into() })?;

    let mut opts = Options::default();
    opts.create_if_missing(false);
    opts.create_missing_column_families(false);
    let db = DB::open_cf_descriptors(&opts, &db_path, column_family_descriptors())?;
    let cf = db
        .cf_handle(CF_REPOSITORY_VIEW)
        .ok_or("missing repository_view column family")?;

    let mut kv_count = 0u64;
    let mut key_bytes_total = 0u64;
    let mut value_bytes_total = 0u64;
    let mut max_key_bytes = 0usize;
    let mut max_value_bytes = 0usize;

    for entry in db.iterator_cf(cf, IteratorMode::Start) {
        let (key, value) = entry?;
        kv_count += 1;
        key_bytes_total += key.len() as u64;
        value_bytes_total += value.len() as u64;
        max_key_bytes = max_key_bytes.max(key.len());
        max_value_bytes = max_value_bytes.max(value.len());
    }

    let logical_total_bytes = key_bytes_total + value_bytes_total;
    let avg_key_bytes = if kv_count > 0 {
        key_bytes_total as f64 / kv_count as f64
    } else {
        0.0
    };
    let avg_value_bytes = if kv_count > 0 {
        value_bytes_total as f64 / kv_count as f64
    } else {
        0.0
    };

    let out = serde_json::json!({
        "db_path": db_path.display().to_string(),
        "column_family": CF_REPOSITORY_VIEW,
        "kv_count": kv_count,
        "key_bytes_total": key_bytes_total,
        "value_bytes_total": value_bytes_total,
        "logical_total_bytes": logical_total_bytes,
        "db_dir_on_disk_bytes": dir_size(&db_path)?,
        "avg_key_bytes": avg_key_bytes,
        "avg_value_bytes": avg_value_bytes,
        "max_key_bytes": max_key_bytes,
        "max_value_bytes": max_value_bytes,
    });
    println!("{}", serde_json::to_string_pretty(&out)?);
    Ok(())
}