use crate::{Quota, quota_id::QuotaId};

const VACANT: usize = usize::MAX;

pub struct QuotaPoolBuffer {
    pub(crate) quotas: Vec<Quota>,
    pub(crate) last_access_times_ns: Vec<u64>,
    pub(crate) elapsed_since_refills_ns: Vec<u64>,
    generations: Vec<u32>,
    free_list: Vec<usize>,
    slot_to_dense: Vec<usize>,
    dense_to_slot: Vec<usize>,
}

impl QuotaPoolBuffer {
    pub fn new() -> Self {
        Self {
            quotas: Vec::new(),
            last_access_times_ns: Vec::new(),
            elapsed_since_refills_ns: Vec::new(),
            generations: Vec::new(),
            free_list: Vec::new(),
            slot_to_dense: Vec::new(),
            dense_to_slot: Vec::new(),
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            quotas: Vec::with_capacity(capacity),
            last_access_times_ns: Vec::with_capacity(capacity),
            elapsed_since_refills_ns: Vec::with_capacity(capacity),
            generations: Vec::with_capacity(capacity),
            free_list: Vec::with_capacity(capacity),
            slot_to_dense: Vec::with_capacity(capacity),
            dense_to_slot: Vec::with_capacity(capacity),
        }
    }

    pub fn push(&mut self, quota: Quota) -> QuotaId {
        let dense = self.quotas.len();
        let slot = match self.free_list.pop() {
            Some(slot) => {
                self.slot_to_dense[slot] = dense;
                slot
            }
            None => {
                let slot = self.generations.len();
                self.generations.push(0);
                self.slot_to_dense.push(dense);
                slot
            }
        };

        self.quotas.push(quota);
        self.last_access_times_ns.push(0);
        self.elapsed_since_refills_ns.push(0);
        self.dense_to_slot.push(slot);

        QuotaId::new(slot, self.generations[slot])
    }

    pub fn raw_quota(&self, id: QuotaId) -> Option<&Quota> {
        let dense = self.resolve_dense(id)?;
        self.quotas.get(dense)
    }

    pub fn last_access_time_ns(&self, id: QuotaId) -> Option<&u64> {
        let dense = self.resolve_dense(id)?;
        self.last_access_times_ns.get(dense)
    }

    pub fn last_access_time_ns_mut(&mut self, id: QuotaId) -> Option<&mut u64> {
        let dense = self.resolve_dense(id)?;
        self.last_access_times_ns.get_mut(dense)
    }

    pub fn elapsed_since_refill_ns(&self, id: QuotaId) -> Option<&u64> {
        let dense = self.resolve_dense(id)?;
        self.elapsed_since_refills_ns.get(dense)
    }

    pub fn elapsed_since_refill_ns_mut(&mut self, id: QuotaId) -> Option<&mut u64> {
        let dense = self.resolve_dense(id)?;
        self.elapsed_since_refills_ns.get_mut(dense)
    }

    pub fn get_mut(&mut self, id: QuotaId) -> Option<&mut Quota> {
        let dense = self.resolve_dense(id)?;
        self.quotas.get_mut(dense)
    }

    pub fn replace(&mut self, id: QuotaId, quota: Quota) -> Option<Quota> {
        let dense = self.resolve_dense(id)?;
        Some(std::mem::replace(&mut self.quotas[dense], quota))
    }

    pub fn remove(&mut self, id: QuotaId) -> bool {
        self.take(id).is_some()
    }

    pub fn take(&mut self, id: QuotaId) -> Option<Quota> {
        let Some(dense) = self.resolve_dense(id) else {
            return None;
        };

        self.generations[id.index] = self.generations[id.index].wrapping_add(1);
        self.slot_to_dense[id.index] = VACANT;
        self.free_list.push(id.index);

        let quota = self.quotas.swap_remove(dense);
        self.last_access_times_ns.swap_remove(dense);
        self.elapsed_since_refills_ns.swap_remove(dense);

        let removed_slot = self.dense_to_slot.swap_remove(dense);
        debug_assert_eq!(removed_slot, id.index);

        if dense < self.dense_to_slot.len() {
            let moved_slot = self.dense_to_slot[dense];
            self.slot_to_dense[moved_slot] = dense;
        }

        Some(quota)
    }

    fn resolve_dense(&self, id: QuotaId) -> Option<usize> {
        let generation = *self.generations.get(id.index)?;
        if generation != id.generation {
            return None;
        }

        let dense = *self.slot_to_dense.get(id.index)?;
        if dense == VACANT || dense >= self.quotas.len() {
            return None;
        }

        debug_assert_eq!(self.dense_to_slot[dense], id.index);

        Some(dense)
    }
}

#[cfg(test)]
mod tests;
