use std::collections::BTreeMap;
use std::path::PathBuf;

use rocksdb::{DB, IteratorMode, Options};
use rpki::storage::{
    ALL_COLUMN_FAMILY_NAMES, CF_AUDIT_RULE_INDEX, CF_RAW_BY_HASH, CF_REPOSITORY_VIEW,
    CF_RRDP_SOURCE, CF_RRDP_SOURCE_MEMBER, CF_RRDP_URI_OWNER, CF_VCIR,
    column_family_descriptors,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum DbStatsMode {
    Estimate,
    Exact,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum CfGroup {
    CurrentRepositoryView,
    CurrentValidationState,
    CurrentRrdpState,
    LegacyCompatibility,
}

impl CfGroup {
    fn as_str(self) -> &'static str {
        match self {
            Self::CurrentRepositoryView => "current_repository_view",
            Self::CurrentValidationState => "current_validation_state",
            Self::CurrentRrdpState => "current_rrdp_state",
            Self::LegacyCompatibility => "legacy_compatibility",
        }
    }
}

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

Options:
  --db <path>     RocksDB directory
  --exact         Iterate to count keys (slower; default uses RocksDB estimates)
  --help          Show this help

Output groups:
  - current_repository_view: repository_view + raw_by_hash
  - current_validation_state: vcir + audit_rule_index
  - current_rrdp_state: rrdp_source + rrdp_source_member + rrdp_uri_owner
"
    )
}

fn estimate_keys(db: &DB, cf_name: &str) -> Result<Option<u64>, Box<dyn std::error::Error>> {
    let cf = db
        .cf_handle(cf_name)
        .ok_or_else(|| format!("missing column family: {cf_name}"))?;
    Ok(db.property_int_value_cf(cf, "rocksdb.estimate-num-keys")?)
}

fn exact_keys(db: &DB, cf_name: &str) -> Result<u64, Box<dyn std::error::Error>> {
    let cf = db
        .cf_handle(cf_name)
        .ok_or_else(|| format!("missing column family: {cf_name}"))?;
    let mode = IteratorMode::Start;
    let mut count = 0u64;
    for res in db.iterator_cf(cf, mode) {
        res?;
        count += 1;
    }
    Ok(count)
}

fn cf_group(cf_name: &str) -> CfGroup {
    match cf_name {
        CF_REPOSITORY_VIEW | CF_RAW_BY_HASH => CfGroup::CurrentRepositoryView,
        CF_VCIR | CF_AUDIT_RULE_INDEX => CfGroup::CurrentValidationState,
        CF_RRDP_SOURCE | CF_RRDP_SOURCE_MEMBER | CF_RRDP_URI_OWNER => {
            CfGroup::CurrentRrdpState
        }
        _ => CfGroup::LegacyCompatibility,
    }
}

fn summarize_counts<'a>(
    counts: impl IntoIterator<Item = (&'a str, u64)>,
) -> BTreeMap<CfGroup, u64> {
    let mut grouped = BTreeMap::new();
    for (cf_name, count) in counts {
        *grouped.entry(cf_group(cf_name)).or_insert(0) += count;
    }
    grouped
}

fn mode_label(mode: DbStatsMode) -> &'static str {
    match mode {
        DbStatsMode::Estimate => "estimate",
        DbStatsMode::Exact => "exact",
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let argv: Vec<String> = std::env::args().collect();
    if argv.iter().any(|a| a == "--help" || a == "-h") {
        print!("{}", usage());
        return Ok(());
    }

    let mut db_path: Option<PathBuf> = None;
    let mut mode = DbStatsMode::Estimate;
    let mut i = 1usize;
    while i < argv.len() {
        match argv[i].as_str() {
            "--db" => {
                i += 1;
                let v = argv.get(i).ok_or("--db requires a value")?;
                db_path = Some(PathBuf::from(v));
            }
            "--exact" => mode = DbStatsMode::Exact,
            other => return Err(format!("unknown argument: {other}\n\n{}", usage()).into()),
        }
        i += 1;
    }

    let db_path = db_path.ok_or_else(|| format!("--db is required\n\n{}", usage()))?;

    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())?;

    println!("db={}", db_path.display());
    println!("mode={}", mode_label(mode));

    let mut per_cf = Vec::with_capacity(ALL_COLUMN_FAMILY_NAMES.len());
    let mut total: u64 = 0;
    for &name in ALL_COLUMN_FAMILY_NAMES {
        let n = match mode {
            DbStatsMode::Exact => exact_keys(&db, name)?,
            DbStatsMode::Estimate => estimate_keys(&db, name)?.unwrap_or(0),
        };
        total = total.saturating_add(n);
        per_cf.push((name, n));
        println!("{name}={n}");
    }
    println!("total={total}");

    for (group, count) in summarize_counts(per_cf.iter().copied()) {
        println!("group_{}={count}", group.as_str());
    }

    let live = db.live_files()?;
    println!("sst_files={}", live.len());

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cf_group_classifies_current_and_legacy_keyspaces() {
        assert_eq!(cf_group(CF_REPOSITORY_VIEW), CfGroup::CurrentRepositoryView);
        assert_eq!(cf_group(CF_RAW_BY_HASH), CfGroup::CurrentRepositoryView);
        assert_eq!(cf_group(CF_VCIR), CfGroup::CurrentValidationState);
        assert_eq!(
            cf_group(CF_AUDIT_RULE_INDEX),
            CfGroup::CurrentValidationState
        );
        assert_eq!(cf_group(CF_RRDP_SOURCE), CfGroup::CurrentRrdpState);
        assert_eq!(cf_group(CF_RRDP_URI_OWNER), CfGroup::CurrentRrdpState);
        assert_eq!(cf_group("unknown_legacy"), CfGroup::LegacyCompatibility);
    }

    #[test]
    fn summarize_counts_accumulates_by_group() {
        let grouped = summarize_counts([
            (CF_REPOSITORY_VIEW, 5),
            (CF_RAW_BY_HASH, 7),
            (CF_VCIR, 11),
            (CF_AUDIT_RULE_INDEX, 13),
            (CF_RRDP_SOURCE_MEMBER, 19),
        ]);

        assert_eq!(grouped.get(&CfGroup::CurrentRepositoryView), Some(&12));
        assert_eq!(grouped.get(&CfGroup::CurrentValidationState), Some(&24));
        assert_eq!(grouped.get(&CfGroup::CurrentRrdpState), Some(&19));
        assert_eq!(grouped.get(&CfGroup::LegacyCompatibility), None);
    }

    #[test]
    fn usage_mentions_grouped_output_and_exact_mode() {
        let text = usage();
        assert!(text.contains("--exact"), "{text}");
        assert!(text.contains("current_validation_state"), "{text}");
        assert!(text.contains("current_rrdp_state"), "{text}");
    }
}
