lunar-lib 0.11.2

use std::{
    any::{Any, TypeId},
    collections::HashMap,
};

use sled::{
    IVec, Transactional, Tree,
    transaction::{ConflictableTransactionError, TransactionalTree},
};

use crate::{
    database::{Database, DatabaseEntry, Db, TransactionError},
    id::Id,
};

/// Holds a compare and swap value, where `old` is the expected value of the item in the database, and `new` is the data we want to overwrite it with
pub struct CompareAndSwapValue<T> {
    pub old: Option<IVec>,
    pub new: Option<T>,
}

impl<T> Default for CompareAndSwapValue<T> {
    fn default() -> Self {
        Self {
            old: Default::default(),
            new: Default::default(),
        }
    }
}

impl<T> CompareAndSwapValue<T> {
    #[must_use]
    pub fn new(old: Option<IVec>, new: Option<T>) -> Self {
        Self { old, new }
    }
}

/// Holds a group of compare and swap values for the given entry tree
pub struct TreeCompareAndSwap<T> {
    tree: Tree,
    swaps: HashMap<Id<T>, CompareAndSwapValue<T>>,
}

impl<T: DatabaseEntry> TreeCompareAndSwap<T> {
    #[must_use]
    fn new(db: &Db<T::DbInner>) -> Self {
        Self {
            tree: db.entry_tree::<T>(),
            swaps: HashMap::new(),
        }
    }

    #[must_use]
    pub fn tree(&self) -> &Tree {
        &self.tree
    }

    #[must_use]
    pub fn get(&self, id: &Id<T>) -> Option<&CompareAndSwapValue<T>> {
        self.swaps.get(id)
    }

    #[must_use]
    pub fn get_mut(&mut self, id: &Id<T>) -> Option<&mut CompareAndSwapValue<T>> {
        self.swaps.get_mut(id)
    }

    #[must_use]
    pub fn take_out(&mut self, id: &Id<T>) -> Option<CompareAndSwapValue<T>> {
        self.swaps.remove(id)
    }

    pub fn insert(&mut self, id: Id<T>, old: Option<IVec>, new: Option<T>) {
        self.swaps.insert(id, CompareAndSwapValue { old, new });
    }

    pub fn insert_raw(&mut self, id: Id<T>, cas_value: CompareAndSwapValue<T>) {
        self.swaps.insert(id, cas_value);
    }
}

/// Generic wrapper over [`TreeCompareAndSwap<T>`]
trait GenericCompareAndSwap: Any {
    fn tree(&self) -> Tree;
    fn as_any(&self) -> &dyn Any;
    fn as_any_mut(&mut self) -> &mut dyn Any;
    fn apply(&self, tx_tree: &TransactionalTree) -> Result<(), TransactionError>;
}

impl GenericCompareAndSwap for IndexCompareAndSwap {
    fn tree(&self) -> Tree {
        self.tree.clone()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn apply(&self, tx_tree: &TransactionalTree) -> Result<(), TransactionError> {
        for (k, v) in &self.swaps {
            let db_old = tx_tree.get(k)?;

            if db_old == v.old {
                if let Some(new) = &v.new {
                    tx_tree.insert(k, new)?;
                } else {
                    tx_tree.remove(k)?;
                }
            } else {
                return Err(TransactionError::CompareAndSwapError);
            }
        }
        Ok(())
    }
}

impl<Entry: DatabaseEntry> GenericCompareAndSwap for TreeCompareAndSwap<Entry> {
    fn tree(&self) -> Tree {
        self.tree.clone()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    /// Applies all changes [`Self`] holds to the `tx_tree`
    fn apply(&self, tx_tree: &TransactionalTree) -> Result<(), TransactionError> {
        for (k, v) in &self.swaps {
            let db_old = tx_tree.get(k)?;

            if db_old == v.old {
                if let Some(new) = &v.new {
                    let buffer = Vec::from(Entry::VERSION_NUMBER.to_be_bytes());
                    let buffer: IVec = cbor4ii::serde::to_vec(buffer, new)
                        .expect("Cbor4ii failed to serialize. This cannot happen unless a serializer failed")
                        .into();

                    tx_tree.insert(&**k, buffer)?;
                } else {
                    tx_tree.remove(&**k)?;
                }
            } else {
                return Err(TransactionError::CompareAndSwapError);
            }
        }
        Ok(())
    }
}

/// Defines an entire compare-and-swap transaction
pub struct CompareAndSwapTransaction<CasDb: Database> {
    swaps: HashMap<TypeId, Box<dyn GenericCompareAndSwap>>,
    index_swaps: HashMap<&'static str, IndexCompareAndSwap>,
    database: Db<CasDb>,
}

impl<CasDb: Database> CompareAndSwapTransaction<CasDb> {
    #[must_use]
    pub(super) fn with_db(database: Db<CasDb>) -> Self {
        Self {
            swaps: HashMap::new(),
            index_swaps: HashMap::new(),
            database,
        }
    }

    pub(super) fn db(&self) -> &Db<CasDb> {
        &self.database
    }

    /// Returns an immutable reference to an entries compare and swap tree
    pub fn get_request<Entry>(&self) -> Option<&TreeCompareAndSwap<Entry>>
    where
        Entry: DatabaseEntry<DbInner = CasDb>,
    {
        self.swaps.get(&TypeId::of::<Entry>()).map(|boxed| {
            boxed
                .as_any()
                .downcast_ref::<TreeCompareAndSwap<Entry>>()
                .unwrap()
        })
    }

    /// Returns a mutable reference to an entries compare and swap tree
    pub fn get_request_mut<Entry>(&mut self) -> Option<&mut TreeCompareAndSwap<Entry>>
    where
        Entry: DatabaseEntry<DbInner = CasDb>,
    {
        self.swaps.get_mut(&TypeId::of::<Entry>()).map(|boxed| {
            boxed
                .as_any_mut()
                .downcast_mut::<TreeCompareAndSwap<Entry>>()
                .unwrap()
        })
    }

    /// Returns a mutable reference to an entries compare and swap tree, opening a new one if one has not already been opened
    pub fn get_or_new_request<Entry>(&mut self) -> &mut TreeCompareAndSwap<Entry>
    where
        Entry: DatabaseEntry<DbInner = CasDb>,
    {
        self.swaps
            .entry(TypeId::of::<Entry>())
            .or_insert_with(|| Box::new(TreeCompareAndSwap::<Entry>::new(&self.database)))
            .as_any_mut()
            .downcast_mut::<TreeCompareAndSwap<Entry>>()
            .unwrap()
    }

    /// Returns a mutable reference to an entries compare and swap tree, opening a new one if one has not already been opened
    pub fn get_or_new_index(&mut self, name: &'static str) -> &mut IndexCompareAndSwap {
        self.index_swaps.entry(name).or_insert_with(|| {
            let tree = self.database.tree(format!("index_{name}"));
            IndexCompareAndSwap::new(tree)
        })
    }

    /// Applies a [`CompareAndSwapTransaction`] atomically to the database
    ///
    /// # Errors
    ///
    /// This function will error if [`sled`] fails get, insert, or remove a key OR abort with a [`CompareAndSwapError`] if the current value does not match the expected value
    pub(super) fn apply(self, flush: bool) -> Result<(), TransactionError> {
        if self.swaps.is_empty() {
            return Ok(());
        }

        let index_ops = self
            .index_swaps
            .into_values()
            .map(|cas| Box::new(cas) as Box<dyn GenericCompareAndSwap>);

        let (trees, swaps): (Vec<Tree>, Vec<Box<dyn GenericCompareAndSwap>>) = self
            .swaps
            .into_values()
            .chain(index_ops)
            .map(|cas| (cas.tree(), cas))
            .unzip();

        trees.transaction(|tx_trees| {
            for (tree, cas) in tx_trees.iter().zip(swaps.iter()) {
                cas.apply(tree)
                    .map_err(ConflictableTransactionError::Abort)?;
                if flush {
                    tree.flush();
                }
            }
            Ok(())
        })?;
        Ok(())
    }
}

/// Holds a group of compare and swap values for the given index tree
pub struct IndexCompareAndSwap {
    tree: Tree,
    swaps: HashMap<IVec, CompareAndSwapValue<IVec>>,
}

impl IndexCompareAndSwap {
    fn new(tree: Tree) -> Self {
        Self {
            tree,
            swaps: HashMap::new(),
        }
    }

    #[must_use]
    fn take_out(&mut self, key: &IVec) -> Option<CompareAndSwapValue<IVec>> {
        self.swaps.remove(key)
    }

    fn insert(
        &mut self,
        key: impl Into<IVec>,
        old: Option<impl Into<IVec>>,
        new: Option<impl Into<IVec>>,
    ) {
        let old = old.map(Into::into);
        let new = new.map(Into::into);
        self.swaps
            .insert(key.into(), CompareAndSwapValue { old, new });
    }

    fn insert_raw(&mut self, key: impl Into<IVec>, cas_value: CompareAndSwapValue<IVec>) {
        self.swaps.insert(key.into(), cas_value);
    }

    /// Fetches a value from the index, updates it using the provided closure, and returns the old value
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn fetch_and_update(
        &mut self,
        key: impl Into<IVec>,
        f: impl FnOnce(Option<IVec>) -> Option<IVec>,
    ) -> Result<Option<IVec>, TransactionError> {
        let key = key.into();
        let (current_state, mut cas_value) = {
            if let Some(v) = self.take_out(&key) {
                (v.new.clone(), v)
            } else {
                let old = self.tree.get(&key)?;
                (old.clone(), CompareAndSwapValue::new(old, None))
            }
        };

        let old = cas_value.new.clone().or_else(|| cas_value.old.clone());
        cas_value.new = f(current_state).clone();
        self.swaps.insert(key, cas_value);
        Ok(old)
    }

    pub fn check(&mut self, key: impl AsRef<[u8]>) -> Result<bool, TransactionError> {
        let exists = self.tree.contains_key(key)?;
        Ok(exists)
    }

    /// Upserts a value to the index
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn upsert(
        &mut self,
        key: impl Into<IVec>,
        value: impl Into<IVec>,
    ) -> Result<(), TransactionError> {
        let key = key.into();
        let value = value.into();
        let cas_value = self.take_out(&key);

        if let Some(mut cas_value) = cas_value {
            cas_value.new = Some(value);
            self.insert_raw(key, cas_value);
        } else {
            let old = self.tree.get(&key)?;
            self.insert(key, old, Some(value));
        }
        Ok(())
    }

    /// Deletes a value from the index and returns it
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn delete(&mut self, key: impl Into<IVec>) -> Result<Option<IVec>, TransactionError> {
        let key = &*key.into();
        self.fetch_and_update(key, |_| None)
    }
}