lean_lmdb

use crate::atoms;
use crate::crud::{
    CrudError, copy_bytes, decode_binary, decode_crud_database, map_heed_error, validate_key,
};
use crate::lifecycle::{BinaryDatabase, DatabaseResource, EnvironmentResource, EnvironmentState};
use rustler::{Atom, Term};
use std::sync::Arc;

pub(crate) const MAX_BATCH_OPERATIONS: usize = 10_000;
pub(crate) const MAX_BATCH_BYTES: usize = 64 * 1024 * 1024;

pub(crate) enum BatchOperation {
    Put {
        database: BinaryDatabase,
        key: Vec<u8>,
        value: Vec<u8>,
    },
    Delete {
        database: BinaryDatabase,
        key: Vec<u8>,
    },
}

pub(crate) fn write_batch_operations(
    state: &EnvironmentState,
    operations: &[BatchOperation],
) -> Result<(), CrudError> {
    if operations.is_empty() {
        return Ok(());
    }
    if state.options.read_only {
        return Err(CrudError::ReadOnly);
    }

    let mut transaction = state.env.write_txn().map_err(map_heed_error)?;
    for operation in operations {
        match operation {
            BatchOperation::Put {
                database,
                key,
                value,
            } => database
                .put(&mut transaction, key, value)
                .map_err(map_heed_error)?,
            BatchOperation::Delete { database, key } => {
                database
                    .delete(&mut transaction, key)
                    .map_err(map_heed_error)?;
            }
        }
    }
    transaction.commit().map_err(map_heed_error)
}

pub(crate) fn validate_batch_bounds(
    operation_count: usize,
    aggregate_bytes: usize,
) -> Result<(), CrudError> {
    if operation_count > MAX_BATCH_OPERATIONS || aggregate_bytes > MAX_BATCH_BYTES {
        Err(CrudError::BatchTooLarge)
    } else {
        Ok(())
    }
}

pub(crate) fn decode_batch_operations(
    environment: &EnvironmentResource,
    operations: Term<'_>,
) -> Result<Vec<BatchOperation>, CrudError> {
    if !operations.is_list() {
        return Err(CrudError::InvalidBatch);
    }

    let mut decoded = Vec::new();
    let mut aggregate_bytes = 0_usize;
    let mut remainder = operations;
    while !remainder.is_empty_list() {
        if decoded.len() == MAX_BATCH_OPERATIONS {
            return Err(CrudError::BatchTooLarge);
        }
        let (operation, tail) = remainder
            .list_get_cell()
            .map_err(|_| CrudError::InvalidBatch)?;
        remainder = tail;
        decoded.try_reserve(1).map_err(|_| CrudError::OutOfMemory)?;

        if let Ok((tag, database, key, value)) = operation.decode::<(Atom, Term, Term, Term)>() {
            if tag != atoms::put() {
                return Err(CrudError::InvalidBatch);
            }
            let database = decode_crud_database(database).map_err(|_| CrudError::InvalidBatch)?;
            if !Arc::ptr_eq(&database.state, &environment.state) {
                return Err(CrudError::MixedEnvironments);
            }
            let key = decode_binary(key, CrudError::InvalidBatch)?;
            let value = decode_binary(value, CrudError::InvalidBatch)?;
            validate_key(&environment.state, key.as_slice())?;
            aggregate_bytes = aggregate_bytes
                .checked_add(key.len())
                .and_then(|bytes| bytes.checked_add(value.len()))
                .ok_or(CrudError::BatchTooLarge)?;
            validate_batch_bounds(decoded.len() + 1, aggregate_bytes)?;
            decoded.push(BatchOperation::Put {
                database: database.database,
                key: copy_bytes(key.as_slice())?,
                value: copy_bytes(value.as_slice())?,
            });
        } else if let Ok((tag, database, key)) = operation.decode::<(Atom, Term, Term)>() {
            if tag != atoms::delete() {
                return Err(CrudError::InvalidBatch);
            }
            let database = decode_crud_database(database).map_err(|_| CrudError::InvalidBatch)?;
            if !Arc::ptr_eq(&database.state, &environment.state) {
                return Err(CrudError::MixedEnvironments);
            }
            let key = decode_binary(key, CrudError::InvalidBatch)?;
            validate_key(&environment.state, key.as_slice())?;
            aggregate_bytes = aggregate_bytes
                .checked_add(key.len())
                .ok_or(CrudError::BatchTooLarge)?;
            validate_batch_bounds(decoded.len() + 1, aggregate_bytes)?;
            decoded.push(BatchOperation::Delete {
                database: database.database,
                key: copy_bytes(key.as_slice())?,
            });
        } else {
            return Err(CrudError::InvalidBatch);
        }
    }

    Ok(decoded)
}

pub(crate) enum ExpectedValue<'a> {
    Missing,
    Value(&'a [u8]),
}

pub(crate) enum ReplacementValue<'a> {
    Delete,
    Put(&'a [u8]),
}

#[derive(Debug, Eq, PartialEq)]
pub(crate) enum CompareExchangeResult<T> {
    Exchanged,
    Conflict(Option<T>),
}

pub(crate) fn compare_exchange_binary<T>(
    database: &DatabaseResource,
    key: &[u8],
    expected: &ExpectedValue<'_>,
    replacement: &ReplacementValue<'_>,
    copy_current: impl FnOnce(&[u8]) -> Result<T, CrudError>,
) -> Result<CompareExchangeResult<T>, CrudError> {
    validate_key(&database.state, key)?;
    if database.state.options.read_only {
        return Err(CrudError::ReadOnly);
    }

    let mut transaction = database.state.env.write_txn().map_err(map_heed_error)?;
    let current = database
        .database
        .get(&transaction, key)
        .map_err(map_heed_error)?;
    let matches = match (expected, current) {
        (ExpectedValue::Missing, None) => true,
        (ExpectedValue::Value(expected), Some(current)) => *expected == current,
        _ => false,
    };
    if !matches {
        // LMDB-owned bytes are copied while the transaction is still alive.
        return Ok(CompareExchangeResult::Conflict(
            current.map(copy_current).transpose()?,
        ));
    }

    match replacement {
        ReplacementValue::Delete => {
            database
                .database
                .delete(&mut transaction, key)
                .map_err(map_heed_error)?;
        }
        ReplacementValue::Put(value) => database
            .database
            .put(&mut transaction, key, *value)
            .map_err(map_heed_error)?,
    }
    transaction.commit().map_err(map_heed_error)?;
    Ok(CompareExchangeResult::Exchanged)
}