use miette::{IntoDiagnostic, Result};
use std::cell::RefCell;
use std::collections::HashSet;

use super::{Db, registry};

impl Db {
    pub fn train_dict(&self, ks_name: &str) -> Result<()> {
        let Some((name, cfg)) = registry::trainable()
            .into_iter()
            .find(|(name, _)| *name == ks_name)
        else {
            miette::bail!("unknown keyspace for training: {ks_name}");
        };
        let ks = self
            .keyspace_by_name(name)
            .expect("trainable keyspace is registered");

        tracing::info!("collecting training samples for keyspace {ks_name}...");

        let mut seen_keys = HashSet::new();
        let captured_keys = RefCell::new(Vec::new());
        let progress = RefCell::new((std::time::Instant::now(), 0));
        let samples = ks
            .sample_data_blocks(cfg.samples, |first, last| {
                let passes = !seen_keys.contains(&(first.to_vec(), last.to_vec()));
                if passes {
                    captured_keys
                        .borrow_mut()
                        .push((first.to_vec(), last.to_vec()));
                }

                let mut p = progress.borrow_mut();
                if passes {
                    p.1 += 1;
                }
                if p.0.elapsed() >= std::time::Duration::from_secs(5) {
                    tracing::info!(
                        "collected {}/{} samples for {}...",
                        p.1,
                        cfg.samples,
                        ks_name
                    );
                    p.0 = std::time::Instant::now();
                }

                passes
            })
            .into_diagnostic()?
            .into_iter()
            .zip(captured_keys.into_inner())
            .filter_map(|(sample, keys)| seen_keys.insert(keys).then_some(sample.to_vec()))
            .collect::<Vec<_>>();

        if samples.is_empty() {
            miette::bail!("no samples found for keyspace {ks_name}, skipping training");
        }

        tracing::info!(
            "training zstd dictionary for keyspace {ks_name} ({} samples, {} bytes limit)...",
            samples.len(),
            cfg.dict_size,
        );
        let dict_bytes =
            lsm_tree::train_zstd_dict(&samples, cfg.dict_size as usize).into_diagnostic()?;
        let path = self.path.join(format!("dict_{ks_name}.bin"));
        std::fs::write(&path, &dict_bytes).into_diagnostic()?;
        tracing::info!(
            "saved zstd dictionary for keyspace {ks_name} ({} bytes) to {}. please restart the indexer to use the new dictionary.",
            dict_bytes.len(),
            path.display()
        );

        Ok(())
    }
}