use crate::constants::{APP_DB_DATA_DIR, DB_COLUMNS};
use rocksdb::{ColumnFamilyDescriptor, IteratorMode, OptimisticTransactionDB, Options, ReadOptions};
use serde::{Serialize};
use serde::de::DeserializeOwned;
use crate::models::RocksColumn;

pub struct RocksDB {
    db: OptimisticTransactionDB,
}

impl RocksDB {
    pub fn new(db_opts: Options, cf_opts: Options) -> color_eyre::Result<Self> {
        let cfs = DB_COLUMNS.iter()
            .map(|cf| ColumnFamilyDescriptor::new(cf.to_string(), cf_opts.clone()))
            .collect::<Vec<_>>();
        let db = OptimisticTransactionDB::open_cf_descriptors(
            &db_opts,
            APP_DB_DATA_DIR.as_path(),
            cfs
        )?;
        let rocks = Self {
            db
        };
        Ok(rocks)
    }

    pub fn get_value<TValue, TColumn>(&self, id: TColumn::Id) -> color_eyre::Result<Option<TValue>>
        where TColumn: RocksColumn, TValue: DeserializeOwned
    {
        let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
        let query_res = self.db.get_cf(&cf, serde_json::to_string(&id)?)?;
        if query_res.is_none() {
            return Ok(None);
        }
        Ok(Some(serde_json::from_slice(&query_res.unwrap())?))
    }

    pub fn set_value<TColumn>(&self, value: &TColumn) -> color_eyre::Result<()>
        where TColumn: RocksColumn + Serialize
    {
        let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
        self.db.put_cf(&cf, serde_json::to_string(&value.get_id())?, serde_json::to_string(value)?)?;
        Ok(())
    }

    pub fn get_values<TColumn>(&self, ids: &[TColumn::Id]) -> color_eyre::Result<Vec<TColumn>>
        where TColumn: RocksColumn + DeserializeOwned
    {
        let transaction = self.db.transaction();
        let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
        let mut values = Vec::new();
        for id in ids {
            let query_res = transaction.get_cf(&cf, serde_json::to_string(id)?)?;
            if let Some(res) = query_res {
                let value = serde_json::from_slice(&res)?;
                values.push(value);
            }
        }
        Ok(values)
    }

    pub fn get_all_values<TColumn>(&self) -> color_eyre::Result<Vec<(TColumn::Id, TColumn)>>
        where TColumn: RocksColumn + DeserializeOwned
    {
        let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
        let mut options = ReadOptions::default();
        options.set_async_io(true);
        let values = self.db.iterator_cf_opt(&cf, options, IteratorMode::Start)
            .filter_map(Result::ok)
            .map(|(k, v)|
                (
                    serde_json::from_slice::<TColumn::Id>(&k).unwrap(),
                    serde_json::from_slice::<TColumn>(&v).unwrap()
                )
            )
            .collect::<Vec<_>>();
        Ok(values)
    }


    pub fn set_values<TColumn>(&self, values: &[TColumn]) -> color_eyre::Result<()>
        where TColumn: RocksColumn + Serialize
    {
        let transaction = self.db.transaction();
        let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
        for value in values {
            transaction.put_cf(&cf, serde_json::to_string(&value.get_id())?, serde_json::to_string(value)?)?;
        }
        transaction.commit()?;
        Ok(())
    }
}