//! Rocks DB implementation of the storage layer
use crate::cursor::Cursor;
use nostr_rs_storage_base::{Error, Index, SecondaryIndex, Storage};
use nostr_rs_types::types::{Event, Filter};
use rocksdb::{
    BoundColumnFamily, ColumnFamilyDescriptor, IteratorMode, Options, SliceTransform, WriteBatch,
    DB,
};
use std::{collections::VecDeque, ops::Deref, path::Path, sync::Arc};

mod cursor;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// Internal index name
pub enum ReferenceType {
    Events,
    Author,
    Tags,
    Kind,
    LocalEvents,
    Stream,
}

impl ReferenceType {
    #[inline]
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Events => "events",
            Self::Author => "authors",
            Self::Tags => "tags",
            Self::Kind => "kinds",
            Self::LocalEvents => "local",
            Self::Stream => "stream",
        }
    }
}

/// Storage implementation using RocksDB
pub struct RocksDb {
    db: DB,
}

impl RocksDb {
    /// Creates a new instance passing rocksDb options and a path
    pub fn with_custom_db<T: AsRef<Path>>(options: Options, path: T) -> Result<Self, Error> {
        let db = DB::open_cf_descriptors(
            &options,
            path,
            vec![
                ColumnFamilyDescriptor::new(ReferenceType::Events.as_str(), options.clone()),
                ColumnFamilyDescriptor::new(ReferenceType::Author.as_str(), options.clone()),
                ColumnFamilyDescriptor::new(ReferenceType::Tags.as_str(), options.clone()),
                ColumnFamilyDescriptor::new(ReferenceType::Kind.as_str(), options.clone()),
                ColumnFamilyDescriptor::new(ReferenceType::LocalEvents.as_str(), options.clone()),
                ColumnFamilyDescriptor::new(ReferenceType::Stream.as_str(), options.clone()),
            ],
        )
        .map_err(|e| Error::Internal(e.to_string()))?;
        Ok(Self { db })
    }

    /// Creates a new instance, passing just the path, the default options are
    /// passed
    pub fn new<T: AsRef<Path>>(path: T) -> Result<Self, Error> {
        let mut options = Options::default();
        options.create_if_missing(true);
        options.create_missing_column_families(true);
        options.set_compression_type(rocksdb::DBCompressionType::Snappy);
        options.optimize_for_point_lookup(256 * 1024 * 1024);

        let prefix_extractor = SliceTransform::create_fixed_prefix(1);
        options.set_prefix_extractor(prefix_extractor);

        Self::with_custom_db(options, path)
    }

    #[inline]
    fn reference_to_cf_handle(
        &self,
        namespace: ReferenceType,
    ) -> Result<Arc<BoundColumnFamily>, Error> {
        self.db
            .cf_handle(namespace.as_str())
            .ok_or(Error::Internal("Unknown db-family".to_owned()))
    }
}

#[async_trait::async_trait]
impl Storage for RocksDb {
    type Cursor<'a> = Cursor<'a>;

    async fn get_local_events(&self, limit: Option<usize>) -> Result<Cursor<'_>, Error> {
        let cf_handle = self.reference_to_cf_handle(ReferenceType::LocalEvents)?;
        Ok(Cursor::new(
            self,
            None,
            vec![],
            None,
            Some(self.db.iterator_cf(&cf_handle, IteratorMode::Start)),
            limit,
        ))
    }

    async fn set_local_event(&self, event: &Event) -> Result<(), Error> {
        let event_id = &event.id;
        let secondary_index = SecondaryIndex::new(event_id, event.created_at());
        self.db
            .put_cf(
                &self.reference_to_cf_handle(ReferenceType::LocalEvents)?,
                secondary_index.index_by([]),
                event_id.deref(),
            )
            .map_err(|e| Error::Internal(e.to_string()))?;
        Ok(())
    }

    async fn store(&self, event: &Event) -> Result<bool, Error> {
        let event_id = &event.id;

        if let Ok(Some(_)) = self.db.get_cf(
            &self.reference_to_cf_handle(ReferenceType::Events)?,
            event_id.deref(),
        ) {
            return Ok(false);
        }
        let secondary_index = SecondaryIndex::new(event_id, event.created_at());
        let author_id = secondary_index.index_by(event.author());
        let json = serde_json::to_vec(event)?;
        let kind: u32 = event.kind().into();
        let kind_id = secondary_index.index_by(kind.to_be_bytes());

        let mut buffer = WriteBatch::default();

        buffer.put_cf(
            &self.reference_to_cf_handle(ReferenceType::Events)?,
            event_id.deref(),
            json,
        );
        buffer.put_cf(
            &self.reference_to_cf_handle(ReferenceType::Author)?,
            author_id,
            event_id.deref(),
        );
        buffer.put_cf(
            &self.reference_to_cf_handle(ReferenceType::Kind)?,
            kind_id,
            event_id.deref(),
        );

        buffer.put_cf(
            &self.reference_to_cf_handle(ReferenceType::Stream)?,
            secondary_index.index_by([]),
            event_id.deref(),
        );

        for tag in event.tags().iter() {
            if let Some(tag) = tag.clone().into_bytes() {
                let tag_id = secondary_index.index_by(&tag);
                buffer.put_cf(
                    &self.reference_to_cf_handle(ReferenceType::Tags)?,
                    tag_id,
                    event_id.deref(),
                );
            }
        }

        self.db
            .write(buffer)
            .map_err(|e| Error::Internal(e.to_string()))?;

        Ok(true)
    }

    async fn get_event<T: AsRef<[u8]> + Send + Sync>(&self, id: T) -> Result<Option<Event>, Error> {
        let id = id.as_ref();
        if let Some(value) = self
            .db
            .prefix_iterator_cf(&self.reference_to_cf_handle(ReferenceType::Events)?, id)
            .next()
        {
            let (key, value) = value.map_err(|e| Error::Internal(e.to_string()))?;

            if key.starts_with(id) {
                return Ok(Some(serde_json::from_slice(&value)?));
            }
        }

        Ok(None)
    }

    async fn get_by_filter(&self, mut query: Filter) -> Result<Cursor<'_>, Error> {
        let limit = if query.limit == 0 {
            None
        } else {
            Some(query.limit.try_into()?)
        };

        let query_index: Index = (&mut query).into();

        let (index, secondary_index_iterator, prefixes) = match query_index {
            Index::Tag(tag_name, tags) => (
                Some(ReferenceType::Tags),
                None,
                tags.into_iter()
                    .map(|tag| tag.into_bytes(&tag_name))
                    .collect::<VecDeque<_>>(),
            ),
            Index::Id(ids) => (
                None,
                None,
                ids.into_iter()
                    .map(|id| id.to_vec())
                    .collect::<VecDeque<_>>(),
            ),
            Index::Author(authors) => (
                Some(ReferenceType::Author),
                None,
                authors
                    .into_iter()
                    .map(|author| author.to_vec())
                    .collect::<VecDeque<_>>(),
            ),
            Index::Kind(kinds) => (
                Some(ReferenceType::Kind),
                None,
                kinds
                    .into_iter()
                    .map(|kind| kind.into_bytes())
                    .collect::<VecDeque<_>>(),
            ),
            Index::TableScan => (
                Some(ReferenceType::Stream),
                Some(self.db.iterator_cf(
                    &self.reference_to_cf_handle(ReferenceType::Stream)?,
                    IteratorMode::Start,
                )),
                VecDeque::new(),
            ),
        };

        Ok(Cursor::new(
            self,
            index,
            prefixes.into(),
            Some(query.into()),
            secondary_index_iterator,
            limit,
        ))
    }
}

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

    async fn new_instance(path: &str) -> RocksDb {
        RocksDb::new(path).expect("valid db")
    }

    async fn destroy_instance(path: &str) {
        std::fs::remove_dir_all(path).expect("deleted file");
    }

    nostr_rs_storage_base::storage_test!(RocksDb, new_instance, destroy_instance);
}
