use std::sync::Arc;

use ahash::AHashMap;
use common::ext::BitSliceExt as _;
use common::is_alive_lock::IsAliveLock;
use common::mmap::MmapBitSlice;
use parking_lot::RwLock;

use crate::common::Flusher;
use crate::common::operation_error::OperationError;

/// A wrapper around `MmapBitSlice` that delays writing changes to the underlying file until they get
/// flushed manually.
/// This expects the underlying MmapBitSlice not to grow in size.
#[derive(Debug)]
pub struct MmapBitSliceBufferedUpdateWrapper {
    bitslice: Arc<RwLock<MmapBitSlice>>,
    len: usize,
    pending_updates: Arc<RwLock<AHashMap<usize, bool>>>,
    /// Lock to prevent concurrent flush and drop
    is_alive_flush_lock: IsAliveLock,
}

impl MmapBitSliceBufferedUpdateWrapper {
    pub fn new(bitslice: MmapBitSlice) -> Self {
        let len = bitslice.len();
        Self {
            bitslice: Arc::new(RwLock::new(bitslice)),
            len,
            pending_updates: Arc::new(RwLock::new(AHashMap::new())),
            is_alive_flush_lock: IsAliveLock::new(),
        }
    }

    /// Sets the bit at `index` to `value` buffered.
    ///
    /// ## Panics
    /// Panics if the index is out of bounds.
    pub fn set(&self, index: usize, value: bool) {
        assert!(index < self.len, "index {index} out of range: {}", self.len);
        self.pending_updates.write().insert(index, value);
    }

    pub fn get(&self, index: usize) -> Option<bool> {
        if index >= self.len {
            return None;
        }
        if let Some(value) = self.pending_updates.read().get(&index) {
            Some(*value)
        } else {
            self.bitslice.read().get_bit(index)
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Removes from `pending_updates` all results that are flushed.
    /// If values in `pending_updates` are changed, do not remove them.
    fn reconcile_persisted_updates(
        pending_updates: &RwLock<AHashMap<usize, bool>>,
        persisted: AHashMap<usize, bool>,
    ) {
        pending_updates
            .write()
            .retain(|point_id, a| persisted.get(point_id).is_none_or(|b| a != b));
    }

    pub fn flusher(&self) -> Flusher {
        let updates = {
            let updates_guard = self.pending_updates.read();
            if updates_guard.is_empty() {
                return Box::new(|| Ok(()));
            }
            updates_guard.clone()
        };

        let bitslice = Arc::downgrade(&self.bitslice);
        let pending_updates_weak = Arc::downgrade(&self.pending_updates);
        let is_alive_flush_lock = self.is_alive_flush_lock.handle();

        Box::new(move || {
            let (Some(is_alive_flush_guard), Some(bitslice), Some(pending_updates_arc)) = (
                is_alive_flush_lock.lock_if_alive(),
                bitslice.upgrade(),
                pending_updates_weak.upgrade(),
            ) else {
                // Already dropped, skip flush
                log::trace!("MmapBitsliceBuffered was dropped, cancelling flush");
                return Err(OperationError::cancelled(
                    "Aborted flushing on a dropped MmapBitSliceBufferedUpdateWrapper instance",
                ));
            };

            let mut mmap_slice_write = bitslice.write();
            for (index, value) in updates.iter() {
                mmap_slice_write.set(*index, *value);
            }
            mmap_slice_write.flusher()()?;

            // Keep the guard till here to prevent concurrent drop/flushes
            // We don't touch files from here on and can drop the alive guard
            drop(is_alive_flush_guard);

            Self::reconcile_persisted_updates(&pending_updates_arc, updates);

            Ok(())
        })
    }
}
