use crate::{
    error::TimeoutError,
    id::{IndexId, RecRef, SegmentId},
    persy::PersyImpl,
};
use std::{collections::HashSet, time::Duration};

#[derive(Clone)]
pub(crate) struct IndexDataLocks {
    locked_index_pages: HashSet<RecRef>,
    index_segments: HashSet<SegmentId>,
    timeout: Duration,
}
impl IndexDataLocks {
    pub(crate) fn new(timeout: Duration) -> Self {
        Self {
            locked_index_pages: Default::default(),
            index_segments: Default::default(),
            timeout,
        }
    }

    pub fn is_record_locked(&self, id: &RecRef) -> bool {
        self.locked_index_pages.contains(id)
    }
    pub fn is_segment_locked(&self, id: &SegmentId) -> bool {
        self.index_segments.contains(id)
    }
    pub fn read_lock_indexes(&mut self, persy_impl: &PersyImpl, to_lock: &[IndexId]) -> Result<(), TimeoutError> {
        let address = persy_impl.address();
        let mut index_segments = Vec::new();
        for index_id in to_lock {
            index_segments.push(index_id.get_meta_id());
            index_segments.push(index_id.get_data_id());
        }
        index_segments.sort();
        address.acquire_segments_read_lock(&index_segments, self.timeout)?;
        for segment in index_segments {
            self.index_segments.insert(segment);
        }
        Ok(())
    }

    pub fn extract(mut self, records: &[RecRef]) -> (HashSet<RecRef>, HashSet<SegmentId>) {
        for rec in records {
            self.locked_index_pages.remove(rec);
        }
        (self.locked_index_pages, self.index_segments)
    }

    pub(crate) fn add_locks(&mut self, locks: &[RecRef]) {
        self.locked_index_pages.extend(locks);
    }
}
