lunar-lib 0.9.0

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

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

use crate::{
    database::{
        Createable, CustomTransactionError, Database, DatabaseEntry, DbHandle, DbKey, Deleteable,
        EntryId, Mergeable, TransactionError, deserialize_from_ivec, serialize_to_ivec,
        sled_get_raw,
    },
    trace,
};

#[derive(Debug, Clone)]
/// 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<Entry: DatabaseEntry> {
    pub old: Option<IVec>,
    pub new: Option<Entry>,
}

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

#[derive(Debug)]
/// Holds a group of compare and swap values for the given tree
pub struct TreeCompareAndSwap<Entry: DatabaseEntry> {
    tree: Tree,
    swaps: HashMap<DbKey, CompareAndSwapValue<Entry>>,
}

impl<T: DatabaseEntry> TreeCompareAndSwap<T> {
    fn new(db: &T::Db) -> Self {
        Self {
            tree: T::__tree(db),
            swaps: HashMap::new(),
        }
    }

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

/// Generic wrapper over [`TreeCompareAndSwap<T>`]
trait GenericCompareAndSwap: Any + std::fmt::Debug {
    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<Entry: DatabaseEntry> GenericCompareAndSwap for TreeCompareAndSwap<Entry> {
    fn tree(&self) -> &Tree {
        &self.tree
    }

    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.as_deref().is_some_and(|v| v[4] != 0) {
                return Err(TransactionError::ReadOnly);
            }

            if db_old == v.old {
                if let Some(new) = &v.new {
                    tx_tree.insert(k, serialize_to_ivec(new))?;
                } else {
                    tx_tree.remove(k)?;
                }
            } else {
                trace!(
                    "CAS mismatch on key {:?}: Expected {:?}. Found {:?}",
                    k, v.old, db_old
                );
                return Err(TransactionError::CompareAndSwapError);
            }
        }
        Ok(())
    }
}

/// Defines an entire compare-and-swap transaction
pub struct CompareAndSwapTransaction<CasDb: Database> {
    swaps: HashMap<TypeId, Box<dyn GenericCompareAndSwap>>,
    database: DbHandle<CasDb>,
}

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

    /// 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<Db = 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()
    }

    /// Gets the latest version of the item in the transaction context, looking it up in the database if its unmodified
    ///
    /// # Errors
    ///
    /// Errors if [`DatabaseEntry::db_get()`] errors
    pub fn tx_get<Id>(&self, id: Id) -> Result<Option<Id::Entry>, TransactionError>
    where
        Id: EntryId<IdDb = CasDb>,
    {
        if let Some(boxed) = self.swaps.get(&TypeId::of::<Id::Entry>()) {
            let cas_tree = boxed
                .as_any()
                .downcast_ref::<TreeCompareAndSwap<Id::Entry>>()
                .unwrap();
            if let Some(get) = cas_tree.swaps.get(id.as_bytes()) {
                return Ok(get.new.clone());
            }
        }

        let tree = Id::Entry::__tree(&*self.database);
        let raw = sled_get_raw(&tree, id.as_bytes())?;
        raw.map(deserialize_from_ivec).transpose()
    }

    /// Returns [`true`] if an item exists in the latest transaction context, else, returns [`false`]
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn tx_check<Id>(&self, id: Id) -> Result<bool, TransactionError>
    where
        Id: EntryId<IdDb = CasDb>,
    {
        if let Some(boxed) = self.swaps.get(&TypeId::of::<Id::Entry>()) {
            let cas_tree = boxed
                .as_any()
                .downcast_ref::<TreeCompareAndSwap<Id::Entry>>()
                .unwrap();
            if let Some(get) = cas_tree.swaps.get(id.as_bytes()) {
                return Ok(get.new.is_some());
            }
        }

        let tree = Id::Entry::__tree(&*self.database);
        Ok(tree.contains_key(id.as_bytes())?)
    }

    /// Returns a batch of items from the database
    ///
    /// # Errors
    ///
    /// - Errors with [`TransactionError::MissingEntry`] if any key does not exist
    /// - Errors if [`CompareAndSwapTransaction::tx_get()`] errors
    pub fn tx_get_batch<Entry, I>(&self, items: I) -> Result<Vec<Entry>, TransactionError>
    where
        Entry: DatabaseEntry<Db = CasDb>,
        I: IntoIterator<Item = Entry::Id>,
    {
        items
            .into_iter()
            .map(|id| self.tx_get(id)?.ok_or(TransactionError::MissingEntry))
            .collect()
    }

    /// Removes an entry from the database
    ///
    /// # Database Safety
    ///
    /// This function will remove an item from the database plainly. It does not care to check if any other items point to it, as it does not know.
    /// Make sure any item you are removing from the database has no references pointing to it
    ///
    /// For a safe version, see [`CompareAndSwapTransaction::tx_delete()`] for types that implement [`Deleteable`]
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn tx_remove<Id>(&mut self, key: Id) -> Result<(), TransactionError>
    where
        Id: EntryId<IdDb = CasDb>,
    {
        let db = self.database.clone();
        let request = self.get_or_new_request::<Id::Entry>();

        let key = *key.as_bytes();
        if let Some(get_mut) = request.swaps.get_mut(&key) {
            get_mut.new = None;
        } else {
            let old = sled_get_raw(&Id::Entry::__tree(&*db), &key)?;
            request
                .swaps
                .insert(key, CompareAndSwapValue { old, new: None });
        }
        Ok(())
    }

    /// Upserts an entry from the database
    ///
    /// # Database Safety
    ///
    /// This function will upsert an item from the database plainly. It does not care to check if it points to other items that don't point back, as it does not know.
    /// Make sure any item you are upserting is fully linked
    ///
    /// For a safe version, see [`CompareAndSwapTransaction::tx_create()`] for types that implement [`Createable`]
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn tx_upsert<Entry: DatabaseEntry<Db = CasDb>>(
        &mut self,
        mut item: Entry,
    ) -> Result<(), TransactionError> {
        let db = self.database.clone();
        item.pre_upsert(self)?;

        let request = self.get_or_new_request::<Entry>();

        let key = *item.id().as_bytes();

        if let Some(get_mut) = request.swaps.get_mut(&key) {
            get_mut.new = Some(item);
        } else {
            let old = sled_get_raw(&Entry::__tree(&*db), &key)?;
            request.swaps.insert(
                key,
                CompareAndSwapValue {
                    old,
                    new: Some(item),
                },
            );
        }
        Ok(())
    }

    /// Inserts an entry from the database
    ///
    /// # Database Safety
    ///
    /// This function will insert an item from the database plainly. It does not care to check if it points to other items that don't point back, as it does not know.
    /// Make sure any item you are upserting is fully linked
    ///
    /// For a safe version, see [`CompareAndSwapTransaction::tx_create()`] for types that implement [`Createable`]
    ///
    /// # Errors
    ///
    /// - Errors with [`TransactionError::AlreadyInDatabase`] if the item already exists
    /// - Errors if [`CompareAndSwapTransaction::tx_upsert()`] errors
    pub fn tx_insert<Entry: DatabaseEntry<Db = CasDb>>(
        &mut self,
        item: Entry,
    ) -> Result<(), TransactionError> {
        if self.tx_get(item.id())?.is_some() {
            return Err(TransactionError::AlreadyInDatabase);
        }

        self.tx_upsert(item)?;
        Ok(())
    }

    pub fn tx_merge<Entry>(&mut self, item: Entry, from: Entry) -> Result<Entry, TransactionError>
    where
        Entry: DatabaseEntry<Db = CasDb> + Mergeable,
    {
        item.tx_merge(from, self)
    }

    /// Attempts to merge the input entry into the existing entry. Upserts if the entry is not found
    ///
    /// # Database Safety
    ///
    /// See [`CompareAndSwapTransaction::tx_upsert()`]
    ///
    /// # Errors
    ///
    /// Errors if `sled` fails to open the entry's tree
    pub fn tx_patch<Entry>(&mut self, item: Entry) -> Result<Entry, TransactionError>
    where
        Entry: DatabaseEntry<Db = CasDb> + Mergeable,
    {
        if let Some(current) = self.tx_get(item.id())? {
            current.tx_merge(item, self)
        } else {
            self.tx_upsert(item.clone())?;
            Ok(item)
        }
    }

    pub fn tx_merge_batch<Entry>(
        &mut self,
        item: Entry,
        from: impl IntoIterator<Item = Entry>,
    ) -> Result<Entry, TransactionError>
    where
        Entry: DatabaseEntry<Db = CasDb> + Mergeable,
    {
        item.tx_merge_batch(from, self)
    }

    pub fn tx_delete<Entry>(&mut self, item: Entry) -> Result<(), TransactionError>
    where
        Entry: DatabaseEntry<Db = CasDb> + Deleteable,
    {
        item.tx_delete(self)
    }

    pub fn tx_delete_id<Id>(&mut self, id: Id) -> Result<(), TransactionError>
    where
        Id: EntryId<IdDb = CasDb, Entry: Deleteable>,
    {
        if let Some(item) = self.tx_get(id)? {
            item.tx_delete(self)
        } else {
            Ok(())
        }
    }

    pub fn tx_create<Entry>(
        &mut self,
        args: <Entry as Createable>::CreateArgs,
    ) -> Result<Entry, CustomTransactionError<<Entry as Createable>::Err>>
    where
        Entry: DatabaseEntry<Db = CasDb> + Createable,
    {
        Entry::tx_create(self, args)
    }
}

/// 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(crate) fn apply_cas_tx<CasDb: Database>(
    tx: CompareAndSwapTransaction<CasDb>,
    flush: bool,
) -> Result<(), TransactionError> {
    if tx.swaps.is_empty() {
        return Ok(());
    }

    let (trees, swaps): (Vec<&Tree>, Vec<&Box<dyn GenericCompareAndSwap>>) =
        tx.swaps.values().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(())
}

/// Calls a closure that returns a [`CompareAndSwapTransaction`] and applies it atomically to the database
///
/// This function will be called again if a [`CompareAndSwapError`] occurs
///
/// # Errors
///
/// This function will error if `f()` returns an error, or if [`apply_cas_tx()`] fails with an error other than [`CompareAndSwapError`]
pub fn db_transaction<CasDb: Database, F, T, E>(
    mut f: F,
    db: DbHandle<CasDb>,
    flush: bool,
) -> Result<T, CustomTransactionError<E>>
where
    F: FnMut(&mut CompareAndSwapTransaction<CasDb>) -> Result<T, CustomTransactionError<E>>,
{
    loop {
        let mut cas_tx = CompareAndSwapTransaction::with_db(db.clone());

        let t = f(&mut cas_tx)?;

        match apply_cas_tx(cas_tx, flush) {
            Ok(()) => return Ok(t),
            Err(TransactionError::CompareAndSwapError) => {
                trace!("Transaction (Not sync) ran into a CAS error and is retrying.");
            }
            Err(err) => return Err(CustomTransactionError::Transaction(err)),
        }
    }
}