rocksolid 2.7.0

use crate::bytes::AsBytes;
use crate::cf_store::RocksDbCFStore;
use crate::error::{StoreError, StoreResult};
use crate::serialization;
use crate::types::ValueWithExpiry;
use crate::MergeValue;

use rocksdb::WriteBatch;
use serde::{de::DeserializeOwned, Serialize};
use std::fmt::Debug;
use std::hash::Hash;
use std::mem::ManuallyDrop;

/// Builds and executes a sequence of write operations atomically on a **single, specified Column Family**.
///
/// Create an instance using `RocksDbCFStore::batch_writer("cf_name")` or `RocksDbStore::batch_writer()`
/// (which defaults to the default Column Family).
///
/// Add operations using methods like `set`, `delete`, `merge`. These operations will
/// implicitly target the Column Family this `BatchWriter` was created for.
///
/// The batch is executed against the database when `.commit()` is called.
/// If the `BatchWriter` is dropped before `.commit()` or `.discard()` is called,
/// a warning will be logged, and the operations will NOT be applied.
pub struct BatchWriter<'a> {
  store: &'a RocksDbCFStore,
  batch: ManuallyDrop<WriteBatch>,
  cf_name: String,
  committed_or_discarded: bool,
}

impl<'a> BatchWriter<'a> {
  pub(crate) fn new(store: &'a RocksDbCFStore, cf_name: String) -> Self {
    BatchWriter {
      store,
      batch: ManuallyDrop::new(WriteBatch::default()),
      cf_name,
      committed_or_discarded: false,
    }
  }

  fn check_not_committed(&self) -> StoreResult<()> {
    if self.committed_or_discarded {
      Err(StoreError::Other(
        "BatchWriter already committed or discarded".to_string(),
      ))
    } else {
      Ok(())
    }
  }

  pub fn set<Key, Val>(&mut self, key: Key, val: &Val) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    Val: Serialize,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let sv = serialization::serialize_value(val)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.put(sk, sv);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.put_cf(&handle, sk, sv);
    }
    Ok(self)
  }

  pub fn set_raw<Key>(&mut self, key: Key, raw_val: &[u8]) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.put(sk, raw_val);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.put_cf(&handle, sk, raw_val);
    }
    Ok(self)
  }

  pub fn set_with_expiry<Key, Val>(&mut self, key: Key, val: &Val, expire_time: u64) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    Val: Serialize + DeserializeOwned + Debug,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let vwe = ValueWithExpiry::from_value(expire_time, val)?;
    let sv_with_ts = vwe.serialize_for_storage();
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.put(sk, sv_with_ts);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.put_cf(&handle, sk, sv_with_ts);
    }
    Ok(self)
  }

  pub fn merge<Key, PatchVal>(&mut self, key: Key, merge_value: &MergeValue<PatchVal>) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    PatchVal: Serialize + Debug,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let smo = serialization::serialize_value(merge_value)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.merge(sk, smo);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.merge_cf(&handle, sk, smo);
    }
    Ok(self)
  }

  pub fn merge_raw<Key>(&mut self, key: Key, raw_merge_op: &[u8]) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.merge(sk, raw_merge_op);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.merge_cf(&handle, sk, raw_merge_op);
    }
    Ok(self)
  }

  pub fn delete<Key>(&mut self, key: Key) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    self.check_not_committed()?;
    let sk = serialization::serialize_key(key)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.delete(sk);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.delete_cf(&handle, sk);
    }
    Ok(self)
  }

  pub fn delete_range<Key>(&mut self, start_key: Key, end_key: Key) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    self.check_not_committed()?;
    let sks = serialization::serialize_key(start_key)?;
    let ske = serialization::serialize_key(end_key)?;
    let current_batch = &mut *self.batch;

    if self.cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      current_batch.delete_range(sks, ske);
    } else {
      let handle = self.store.get_cf_handle(&self.cf_name)?;
      current_batch.delete_range_cf(&handle, sks, ske);
    }
    Ok(self)
  }

  /// Provides mutable access to the underlying `rocksdb::WriteBatch`.
  pub fn raw_batch_mut(&mut self) -> StoreResult<&mut WriteBatch> {
    self.check_not_committed()?;
    Ok(&mut *self.batch)
  }

  /// Commits the accumulated batch operations atomically to the database.
  /// Consumes the `BatchWriter`.
  pub fn commit(mut self) -> StoreResult<()> {
    self.check_not_committed()?;
    // Safely take ownership of the WriteBatch from ManuallyDrop
    let batch_to_commit = unsafe { ManuallyDrop::take(&mut self.batch) };
    self
      .store
      .db_raw()
      .write(batch_to_commit)
      .map_err(StoreError::RocksDb)?;
    self.committed_or_discarded = true;
    Ok(())
  }

  /// Explicitly discards the batch without committing any operations.
  /// Consumes the `BatchWriter`.
  pub fn discard(mut self) {
    // No need to check_not_committed, discard is always safe.
    // We still need to take ownership to ensure Drop doesn't try to log for it.
    let _batch_to_discard = unsafe { ManuallyDrop::take(&mut self.batch) };
    self.committed_or_discarded = true;
    // The taken batch will be dropped here.
  }
}

impl<'a> Drop for BatchWriter<'a> {
  fn drop(&mut self) {
    if !self.committed_or_discarded {
      log::warn!(
                "BatchWriter for DB at '{}' (CF: '{}') dropped without calling commit() or discard(). Batch operations were NOT applied.",
                self.store.path(),
                self.cf_name
            );
      // If not committed/discarded, the inner WriteBatch is dropped normally here.
      // WriteBatch has no rollback, so there is nothing else to do.
    }
  }
}

/// Builds and executes a sequence of write operations atomically **across any number of
/// Column Families** in a single `WriteBatch`.
///
/// Unlike [`BatchWriter`], which is bound to one CF for its lifetime, each write method here
/// takes the target `cf_name` as its first argument, so a single batch can mix writes to
/// several CFs. All operations commit atomically when [`commit`](Self::commit) is called.
///
/// Create an instance with `RocksDbCFStore::batch_writer_multi_cf()` or
/// `RocksDbStore::batch_writer_multi_cf()`.
///
/// If the writer is dropped before `.commit()` or `.discard()`, a warning is logged and no
/// operations are applied.
pub struct MultiCfBatchWriter<'a> {
  store: &'a RocksDbCFStore,
  batch: ManuallyDrop<WriteBatch>,
  committed_or_discarded: bool,
}

impl<'a> MultiCfBatchWriter<'a> {
  pub(crate) fn new(store: &'a RocksDbCFStore) -> Self {
    MultiCfBatchWriter {
      store,
      batch: ManuallyDrop::new(WriteBatch::default()),
      committed_or_discarded: false,
    }
  }

  fn check_not_committed(&self) -> StoreResult<()> {
    if self.committed_or_discarded {
      Err(StoreError::Other(
        "MultiCfBatchWriter already committed or discarded".to_string(),
      ))
    } else {
      Ok(())
    }
  }

  /// Applies `default_op` to the batch when `cf_name` is the default CF, otherwise resolves
  /// the CF handle and applies `cf_op`. Centralizes the default-vs-CF branch for every method.
  fn dispatch<FDef, FCf>(&mut self, cf_name: &str, default_op: FDef, cf_op: FCf) -> StoreResult<&mut Self>
  where
    FDef: FnOnce(&mut WriteBatch),
    FCf: FnOnce(&mut WriteBatch, &std::sync::Arc<rocksdb::BoundColumnFamily>),
  {
    self.check_not_committed()?;
    if cf_name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
      default_op(&mut self.batch);
    } else {
      let handle = self.store.get_cf_handle(cf_name)?;
      cf_op(&mut self.batch, &handle);
    }
    Ok(self)
  }

  pub fn set_in<Key, Val>(&mut self, cf_name: &str, key: Key, val: &Val) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    Val: Serialize,
  {
    let sk = serialization::serialize_key(key)?;
    let sv = serialization::serialize_value(val)?;
    self.dispatch(
      cf_name,
      |b| b.put(&sk, &sv),
      |b, h| b.put_cf(h, &sk, &sv),
    )
  }

  pub fn set_raw_in<Key>(&mut self, cf_name: &str, key: Key, raw_val: &[u8]) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    let sk = serialization::serialize_key(key)?;
    self.dispatch(
      cf_name,
      |b| b.put(&sk, raw_val),
      |b, h| b.put_cf(h, &sk, raw_val),
    )
  }

  pub fn set_with_expiry_in<Key, Val>(
    &mut self,
    cf_name: &str,
    key: Key,
    val: &Val,
    expire_time: u64,
  ) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    Val: Serialize + DeserializeOwned + Debug,
  {
    let sk = serialization::serialize_key(key)?;
    let vwe = ValueWithExpiry::from_value(expire_time, val)?;
    let sv = vwe.serialize_for_storage();
    self.dispatch(
      cf_name,
      |b| b.put(&sk, &sv),
      |b, h| b.put_cf(h, &sk, &sv),
    )
  }

  pub fn merge_in<Key, PatchVal>(
    &mut self,
    cf_name: &str,
    key: Key,
    merge_value: &MergeValue<PatchVal>,
  ) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
    PatchVal: Serialize + Debug,
  {
    let sk = serialization::serialize_key(key)?;
    let smo = serialization::serialize_value(merge_value)?;
    self.dispatch(
      cf_name,
      |b| b.merge(&sk, &smo),
      |b, h| b.merge_cf(h, &sk, &smo),
    )
  }

  pub fn merge_raw_in<Key>(&mut self, cf_name: &str, key: Key, raw_merge_op: &[u8]) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    let sk = serialization::serialize_key(key)?;
    self.dispatch(
      cf_name,
      |b| b.merge(&sk, raw_merge_op),
      |b, h| b.merge_cf(h, &sk, raw_merge_op),
    )
  }

  pub fn delete_in<Key>(&mut self, cf_name: &str, key: Key) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    let sk = serialization::serialize_key(key)?;
    self.dispatch(cf_name, |b| b.delete(&sk), |b, h| b.delete_cf(h, &sk))
  }

  pub fn delete_range_in<Key>(&mut self, cf_name: &str, start_key: Key, end_key: Key) -> StoreResult<&mut Self>
  where
    Key: AsBytes + Hash + Eq + PartialEq + Debug,
  {
    let sks = serialization::serialize_key(start_key)?;
    let ske = serialization::serialize_key(end_key)?;
    self.dispatch(
      cf_name,
      |b| b.delete_range(&sks, &ske),
      |b, h| b.delete_range_cf(h, &sks, &ske),
    )
  }

  /// Provides mutable access to the underlying `rocksdb::WriteBatch` for operations not
  /// covered by the typed methods.
  pub fn raw_batch_mut(&mut self) -> StoreResult<&mut WriteBatch> {
    self.check_not_committed()?;
    Ok(&mut *self.batch)
  }

  /// Commits the accumulated batch operations atomically to the database. Consumes the writer.
  pub fn commit(mut self) -> StoreResult<()> {
    self.check_not_committed()?;
    let batch_to_commit = unsafe { ManuallyDrop::take(&mut self.batch) };
    self
      .store
      .db_raw()
      .write(batch_to_commit)
      .map_err(StoreError::RocksDb)?;
    self.committed_or_discarded = true;
    Ok(())
  }

  /// Explicitly discards the batch without committing any operations. Consumes the writer.
  pub fn discard(mut self) {
    let _batch_to_discard = unsafe { ManuallyDrop::take(&mut self.batch) };
    self.committed_or_discarded = true;
  }
}

impl<'a> Drop for MultiCfBatchWriter<'a> {
  fn drop(&mut self) {
    if !self.committed_or_discarded {
      log::warn!(
        "MultiCfBatchWriter for DB at '{}' dropped without calling commit() or discard(). Batch operations were NOT applied.",
        self.store.path(),
      );
    }
  }
}