//! rocksdb backend for hubble-sync
//!
//! app provides column families (one point-oriented, one queue), which it's
//! allowed to share:
//! - this engine will only ever read/write keys under its configured prefix
//! - the app MUST NOT ever write its own keys under that prefix

use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use hubble_sync::{StorageBatch, StorageEngine, StorageError};
use rocksdb::{
    BlockBasedOptions, Cache, ColumnFamily, ColumnFamilyDescriptor, DB, IteratorMode, Options,
    ReadOptions, WriteBatch,
};

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("rocksdb: {0}")]
    Rocks(#[from] rocksdb::Error),

    #[error("missing column family: {0:?}. you might need to initialize it")]
    MissingCf(&'static str),
}

impl StorageError for Error {}

/// point-read-optimized column family for hubble-sync to use.
/// it's fine (+ recommended) to share this CF with point-oriented app data, but
/// you MUST ensure that app data NEVER writes keys under `sync_prefix`.
pub const DEFAULT_POINT_CF: &str = "default";

/// queue-optimized column family for hubble-sync to use.
/// it's fine (+ recommended) to share this CF with queue-oriented app data, but
/// you MUST ensure that app data NEVER writes keys under `sync_prefix`.
pub const DEFAULT_QUEUE_CF: &str = "queue";

/// key prefix that hubble-sync's data will be written under, for both column
/// families. if sharing a column family for app data, you MUST NEVER write keys
/// under hubble-sync's prefix.
pub const DEFAULT_SYNC_PREFIX: &[u8] = &[0x00];

/// RocksDB storage engine implementation for hubble-sync
#[derive(Debug, Clone)]
pub struct RocksEngine {
    db: Arc<DB>,
    prefix: &'static [u8],
    point_cf: &'static str,
    queue_cf: &'static str,
}

#[derive(Debug, Default)]
pub struct RocksEngineConfig {
    pub prefix: Option<&'static [u8]>,
    pub point_cf: Option<&'static str>,
    pub queue_cf: Option<&'static str>,
}

type RocksResult<T> = Result<T, Error>;
type RocksScanItem = RocksResult<(Vec<u8>, Vec<u8>)>;

#[derive(Debug)]
pub struct EasySetup {
    pub block_cache_mb: usize,
}

impl Default for EasySetup {
    fn default() -> EasySetup {
        EasySetup {
            block_cache_mb: 512,
        }
    }
}

impl RocksEngine {
    /// simple mode: let the engine set everything up and hand it to you
    pub fn new_easy(
        path: impl AsRef<Path>,
        prefix: &'static [u8],
        setup: &EasySetup,
    ) -> RocksResult<(Self, Arc<DB>)> {
        let cache = Cache::new_lru_cache(setup.block_cache_mb * 2_usize.pow(20));
        let cfs = recommended_cf_descriptors(&cache);
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);
        let db = Arc::new(DB::open_cf_descriptors(&db_opts, path, cfs)?);
        let engine = Self::new(
            db.clone(),
            &RocksEngineConfig {
                prefix: Some(prefix),
                ..Default::default()
            },
        )?;
        Ok((engine, db))
    }

    pub fn new(db: Arc<DB>, config: &RocksEngineConfig) -> RocksResult<Self> {
        let point_cf = config.point_cf.unwrap_or(DEFAULT_POINT_CF);
        db.cf_handle(point_cf).ok_or(Error::MissingCf(point_cf))?;
        let queue_cf = config.queue_cf.unwrap_or(DEFAULT_QUEUE_CF);
        db.cf_handle(queue_cf).ok_or(Error::MissingCf(queue_cf))?;
        Ok(Self {
            db,
            prefix: config.prefix.unwrap_or(DEFAULT_SYNC_PREFIX),
            point_cf,
            queue_cf,
        })
    }

    pub fn scan_cf(
        &self,
        cf: &ColumnFamily,
        scan_prefix: &[u8],
        from_suffix: &[u8],
    ) -> Box<dyn Iterator<Item = RocksScanItem> + '_> {
        let full_prefix = prefixed(self.prefix, scan_prefix);

        let mut start = Vec::with_capacity(full_prefix.len() + from_suffix.len());
        start.extend_from_slice(&full_prefix);
        start.extend_from_slice(from_suffix);

        let mut opts = ReadOptions::default();
        opts.set_iterate_lower_bound(start);
        if let Some(end) = prefix_end_exclusive(&full_prefix) {
            opts.set_iterate_upper_bound(end);
        }

        let iter = self
            .db
            .iterator_cf_opt(&cf, opts, IteratorMode::Start)
            .map(move |r| {
                r.map(|(k, v)| {
                    // STRIPS THE FULL PREFIX!
                    let k = unprefixed(&full_prefix, &k).to_vec();
                    let v = v.into_vec();
                    (k, v)
                })
                .map_err(Error::Rocks)
            });

        Box::new(iter)
    }
}

impl StorageEngine for RocksEngine {
    type Batch = RocksBatch;
    type Error = Error;

    fn batch(&self) -> Self::Batch {
        RocksBatch {
            db: Arc::clone(&self.db),
            inner: WriteBatch::default(),
            point_cf: self.point_cf,
            queue_cf: self.queue_cf,
            prefix: self.prefix,
        }
    }

    fn get(&self, k: &[u8]) -> RocksResult<Option<Vec<u8>>> {
        let cf = self.db.cf_handle(self.point_cf).expect("point cf to exist");
        Ok(self.db.get_cf(&cf, prefixed(self.prefix, k))?)
    }

    fn get_queue(&self, k: &[u8]) -> RocksResult<Option<Vec<u8>>> {
        let cf = self.db.cf_handle(self.queue_cf).expect("queue cf to exist");
        Ok(self.db.get_cf(&cf, prefixed(self.prefix, k))?)
    }

    fn scan_from(
        &self,
        scan_prefix: &[u8],
        from_suffix: &[u8],
    ) -> Box<dyn Iterator<Item = RocksScanItem> + '_> {
        let cf = self.db.cf_handle(self.point_cf).expect("queue cf to exist");
        self.scan_cf(cf, scan_prefix, from_suffix)
    }

    fn scan_from_queue(
        &self,
        scan_prefix: &[u8],
        from_suffix: &[u8],
    ) -> Box<dyn Iterator<Item = RocksScanItem> + '_> {
        let cf = self.db.cf_handle(self.queue_cf).expect("queue cf to exist");
        self.scan_cf(cf, scan_prefix, from_suffix)
    }

    /// rocks engine forces queue CF flushes every few minutes to keep WAL
    /// anchor moving
    fn maintenance(&self) -> RocksResult<Option<Duration>> {
        let cf = self.db.cf_handle(self.queue_cf).expect("queue cf to exist");
        self.db.flush_cf(&cf)?;
        Ok(Some(Duration::from_secs(120)))
    }
}

pub struct RocksBatch {
    db: Arc<DB>,
    inner: WriteBatch,
    prefix: &'static [u8],
    point_cf: &'static str,
    queue_cf: &'static str,
}

impl RocksBatch {
    /// directly access the storage engine's concrete batch
    ///
    /// unlike the trait methods, does not add a hubble-sync prefix to keys!
    pub fn inner_mut(&mut self) -> &mut WriteBatch {
        &mut self.inner
    }
}

impl StorageBatch<Error> for RocksBatch {
    fn commit(self) -> RocksResult<()> {
        self.db.write(self.inner)?;
        Ok(())
    }
    fn put(&mut self, k: &[u8], v: &[u8]) {
        let cf = self.db.cf_handle(self.point_cf).expect("point CF missing");
        let key = prefixed(self.prefix, k);
        self.inner.put_cf(&cf, key, v);
    }
    fn put_queue(&mut self, k: &[u8], v: &[u8]) {
        let cf = self.db.cf_handle(self.queue_cf).expect("queue CF missing");
        let key = prefixed(self.prefix, k);
        self.inner.put_cf(&cf, key, v);
    }
    fn delete(&mut self, k: &[u8]) {
        let cf = self.db.cf_handle(self.point_cf).expect("point CF missing");
        let key = prefixed(self.prefix, k);
        self.inner.delete_cf(&cf, key);
    }
    fn delete_queue(&mut self, k: &[u8]) {
        let cf = self.db.cf_handle(self.queue_cf).expect("queue CF missing");
        let key = prefixed(self.prefix, k);
        self.inner.delete_cf(&cf, key);
    }
}

fn prefixed(prefix: &[u8], k: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(prefix.len() + k.len());
    out.extend_from_slice(prefix);
    out.extend_from_slice(k);
    out
}

fn unprefixed(prefix: &[u8], k: &[u8]) -> Vec<u8> {
    assert!(
        k.starts_with(prefix),
        "engine bug or data corruption, maybe should be an error?"
    );
    k[prefix.len()..].to_vec()
}

/// smallest key strictly greater than every key beginning with `prefix`.
/// special-case: `None` if prefix is [0xFF, 0xFF, ...] (there is no next key)
pub fn prefix_end_exclusive(prefix: &[u8]) -> Option<Vec<u8>> {
    let mut end = prefix.to_vec();
    // try to increment least-significant-incrementable-byte
    while let Some(last) = end.last_mut() {
        if *last < 0xFF {
            *last += 1;
            return Some(end);
        }
        end.pop();
    }
    None
}

pub fn recommended_cf_descriptors(cache: &Cache) -> Vec<ColumnFamilyDescriptor> {
    vec![
        ColumnFamilyDescriptor::new(DEFAULT_POINT_CF, recommended_point_opts(cache)),
        ColumnFamilyDescriptor::new(DEFAULT_QUEUE_CF, recommended_queue_opts(cache)),
    ]
}

pub fn recommended_point_opts(cache: &Cache) -> Options {
    let mut o = Options::default();
    let mut bbo = BlockBasedOptions::default();
    bbo.set_block_cache(cache);
    o.set_block_based_table_factory(&bbo);
    o
}

pub fn recommended_queue_opts(cache: &Cache) -> Options {
    let mut o = Options::default();
    o.set_write_buffer_size(128 * 2_usize.pow(20)); // 128 MiB
    o.set_max_write_buffer_number(4);
    o.set_min_write_buffer_number_to_merge(2); // insert + delete become friends
    o.add_compact_on_deletion_collector_factory(
        10_000, // window size (how many entries to look at)
        4_000,  // at least this many tombstones (redunant with ratio?)
        0.4,    // ratio of tombstone entries exceeding this
    ); // window, num-deletes trigger, deletion ratio
    let mut bbo = BlockBasedOptions::default();
    bbo.set_block_cache(cache);
    o.set_block_based_table_factory(&bbo);
    o
}
