// Copyright 2026 PingCAP, Inc.

use std::{
    collections::HashMap,
    hash::Hash,
    marker::PhantomData,
    sync::{
        RwLock, atomic,
        atomic::{AtomicU64, Ordering},
    },
};

use crate::time::Instant;

#[cfg(test)]
thread_local! {
    static MOCK_NOW_COARSE_SECONDS: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
}

#[cfg(test)]
fn set_now_for_test(now_secs: Option<u64>) {
    set_now_coarse_seconds_for_test(now_secs);
}

#[cfg(test)]
pub(crate) fn set_now_coarse_seconds_for_test(now_secs: Option<u64>) {
    MOCK_NOW_COARSE_SECONDS.set(now_secs);
}

#[inline]
fn now_coarse_seconds() -> u64 {
    #[cfg(test)]
    if let Some(now_secs) = MOCK_NOW_COARSE_SECONDS.get() {
        return now_secs;
    }
    Instant::now_coarse().second().max(0) as u64
}

/// A concurrent fair queue that partitions tasks by group key.
///
/// Each group owns an inner queue created by `create_group_queue`. [`push`]
/// inserts a task into the group identified by `key`, creating the group on
/// demand. [`pop`] chooses the non-empty group with the smallest virtual time
/// and then pops one task from that group's inner queue.
///
/// This queue only selects groups by their current virtual time. It does not
/// charge virtual time in [`pop`]. Callers should charge reserved or completed
/// work through the [`FairnessRecorder`] implementation, usually after a task
/// is selected, polled, or finished.
///
/// Concurrency and correctness assumptions:
///
/// This queue **does not** provide strict linearizability across group
/// membership, virtual time, total length, and inner queue state. Instead, it
/// provides bounded consistency guarantees that are enough for scheduler-style
/// users: tasks should not be lost, the queue should not look empty while
/// counted work remains, and total length should keep covering each task after
/// a successful push until that task is popped.
///
/// One practical consequence is that [`len`] is a progress hint, not a
/// readiness proof. To preserve the length-covering guarantee after a
/// successful [`push`], this queue increments the total length before
/// publishing the task into the inner queue. A task should be considered fully
/// pushed only after [`push`] returns, or after a notification sent by the
/// producer after that return. If one thread is only pushing and another thread
/// observes `len() > 0` before that point, a concurrent [`pop`] may still
/// return `None` while it races with the in-progress [`push`].
///
/// Within this bounded model, the queue provides the following guarantees:
///
/// * After [`push`] returns, the task is counted by [`len`] until it is popped.
///   A [`pop`] call that starts after that return, or after a producer
///   notification sent after that return, should not return `None` as long as
///   at least one counted task has not been popped. A [`pop`] call that already
///   started while a concurrent [`push`] was still in progress is not covered
///   by that task, even if the [`push`] returns before that [`pop`] returns.
/// * [`drain`] is intended for shutdown-style draining. It snapshots group
///   keys, so callers should stop concurrent producers if they need a complete
///   drain.
/// * These guarantees rely on the inner queue type `Q` being thread-safe and
///   providing consistent [`GroupInnerQueue::push`], [`GroupInnerQueue::pop`],
///   and [`GroupInnerQueue::len`] semantics under concurrent access.
///
/// Performance characteristics:
///
/// * With a stable group set, existing-group pushes, pops, and virtual-time
///   updates use the registry read lock. Inner queue operations run while that
///   read lock is held, so they should be short and non-blocking.
/// * Creating or removing groups takes the registry write lock and blocks
///   concurrent pushes, pops, and virtual-time updates while the membership
///   change is in progress. This is expected to be rare.
/// * [`pop`] scans all registered groups to find the minimum virtual time. This
///   is cheap for a small or moderate number of groups, but it is not suitable
///   for very large group sets or high group churn.
///
/// This type fits workloads with relatively few long-lived groups, frequent
/// per-group task operations, and infrequent group joins/removals. For example,
/// it is suitable for choosing among active keyspace queues, but not for using
/// a different group for every short-lived task.
pub struct MultiGroupFairQueue<K, T, Q> {
    /// Group registry that stores each group queue and its virtual time.
    registry: FairRegistry<K, FairGroup<T, Q>>,
    /// Total counted tasks across groups.
    len: atomic::AtomicUsize,
    /// Factory used when a group is created on first push.
    create_group_queue: Box<dyn Fn(K) -> Q + Send + Sync>,
}

impl<K, T, Q> MultiGroupFairQueue<K, T, Q>
where
    K: FairRegistryKey,
    Q: GroupInnerQueue<T> + Send + Sync,
{
    /// Creates an empty queue.
    ///
    /// `max_vt_lag` is passed to the underlying [`FairRegistry`] and uses the
    /// same virtual-time unit as the caller's accounting. `create_group_queue`
    /// is called with the group key when a group is first created.
    pub fn new(
        max_vt_lag: u64,
        create_group_queue: impl Fn(K) -> Q + Send + Sync + 'static,
    ) -> Self {
        Self {
            registry: FairRegistry::new(max_vt_lag),
            len: atomic::AtomicUsize::new(0),
            create_group_queue: Box::new(create_group_queue),
        }
    }

    /// Pushes `task` into the group identified by `key`.
    ///
    /// The group is created on demand. After this method returns, the task is
    /// counted by [`MultiGroupFairQueue::len`] until it is popped or drained.
    /// Creating a new group takes the registry write lock; pushing into an
    /// existing group uses the registry read lock.
    pub fn push(&self, key: K, task: T) {
        self.registry.ensure_entry_with_op(
            key,
            || FairGroup::new(self.create_group_queue.as_ref()(key)),
            |g| {
                // Count before publishing into the group queue. If the task were
                // published first, a racing pop could take it and decrement the
                // total length before this push increments it.
                self.len.fetch_add(1, Ordering::AcqRel);
                g.push(task);
            },
        );
    }

    /// Pops one task from the currently fairest non-empty group.
    ///
    /// Fairness is based on each group's current virtual time. This method does
    /// not charge virtual time; callers should do that through
    /// [`FairnessRecorder`]. `None` means no task was available at the time of
    /// selection, but it may also be observed by a call that started while
    /// racing with an in-progress [`MultiGroupFairQueue::push`] whose task had
    /// not yet been fully published. Once a [`MultiGroupFairQueue::push`]
    /// returns, later [`MultiGroupFairQueue::pop`] calls are expected to
    /// observe that task, unless another pop has already taken it.
    pub fn pop(&self) -> Option<T> {
        if self.is_empty() {
            return None;
        }
        // A group can look non-empty but yield no task in this attempt. This
        // may happen when a concurrent push is between its length increment and
        // task publication, or when another worker pops the last visible task
        // after the length snapshot.
        //
        // Without remembering skipped groups, the next iteration may choose the
        // same group again because its virtual time is still the smallest and
        // its length may still report non-empty. In the worst case, this pop
        // call busy-waits on that one false-ready group until the racing push
        // publishes its task, while ready tasks in other groups are not tried.
        let mut skip_list = vec![];
        loop {
            let popped = self.registry.operate_min_vt_entry(
                |key, g| !g.is_empty() && !skip_list.contains(&key),
                |key, g| g.pop().ok_or(key),
            )?;
            match popped {
                Ok(task) => {
                    self.len.fetch_sub(1, Ordering::AcqRel);
                    return Some(task);
                }
                Err(key) => {
                    // The selected group had no task by the time we popped it.
                    // It may become visible after a racing push returns, or it
                    // may have already been taken by another worker.
                    skip_list.push(key);
                }
            }
        }
    }

    /// Returns the current counted task count.
    ///
    /// The value is a snapshot and may become stale immediately. It is a
    /// progress hint, not a proof that [`MultiGroupFairQueue::pop`] must return
    /// a task.
    #[inline]
    pub fn len(&self) -> usize {
        self.len.load(Ordering::Acquire)
    }

    /// Returns whether the current counted task count is zero.
    ///
    /// This is equivalent to `self.len() == 0` and has the same snapshot
    /// semantics as [`MultiGroupFairQueue::len`].
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns a snapshot of the current group count.
    ///
    /// The returned value may become stale immediately if another thread adds
    /// or removes groups.
    #[inline]
    pub fn group_count(&self) -> usize {
        self.registry.entry_count()
    }

    /// Drains tasks from all groups present in the current key snapshot.
    ///
    /// This is intended for shutdown-style draining. It does not prevent
    /// concurrent producers from adding new tasks, and it may miss groups that
    /// are created after the key snapshot is taken.
    pub fn drain(&self) {
        for key in self.registry.keys() {
            self.registry.operate_entry(key, |group| {
                while group.pop().is_some() {
                    self.len.fetch_sub(1, Ordering::AcqRel);
                }
            });
        }
    }

    /// Removes empty groups that have been inactive for longer than
    /// `inactive_secs`.
    ///
    /// This is a best-effort cleanup. It snapshots keys first, then removes a
    /// group only if it is still empty and inactive under the registry write
    /// lock.
    pub fn remove_inactive_groups(&self, inactive_secs: u64) {
        let is_group_expired = |group: &FairGroup<T, Q>| {
            group.is_empty() && group.secs_since_last_access() > inactive_secs
        };
        for key in self.registry.keys() {
            let inactive = self
                .registry
                .operate_entry(key, is_group_expired)
                .unwrap_or(false);
            if inactive {
                self.registry.remove_if(key, is_group_expired);
            }
        }
    }
}

impl<K, T, Q> FairnessRecorder<K> for MultiGroupFairQueue<K, T, Q>
where
    K: FairRegistryKey,
    Q: GroupInnerQueue<T> + Send + Sync,
{
    #[inline]
    fn increase_vt(&self, k: K, vt: u64) {
        self.registry.increase_vt(k, vt);
    }

    #[inline]
    fn decrease_vt(&self, k: K, vt: u64) {
        self.registry.decrease_vt(k, vt);
    }
}

/// Inner per-group queue used by [`MultiGroupFairQueue`].
///
/// Implementations must be thread-safe when used with [`MultiGroupFairQueue`].
/// After [`GroupInnerQueue::push`] returns, the task must be included in
/// [`GroupInnerQueue::len`] until it is returned by [`GroupInnerQueue::pop`].
/// As long as an unpopped task remains in this inner queue, repeated
/// [`GroupInnerQueue::pop`] calls must be able to return it.
///
/// This trait intentionally exposes [`GroupInnerQueue::len`] but not
/// `is_empty`. [`MultiGroupFairQueue`] derives emptiness from `len() == 0` so
/// implementors cannot provide inconsistent length and emptiness semantics.
#[allow(clippy::len_without_is_empty)]
pub trait GroupInnerQueue<T> {
    /// Pushes a task into this inner queue.
    fn push(&self, t: T);

    /// Pops one task from this inner queue, if one is currently available.
    fn pop(&self) -> Option<T>;

    /// Returns the current task count for this inner queue.
    ///
    /// The value may be a concurrent snapshot, but it must count every task
    /// after a successful [`GroupInnerQueue::push`] until that task is popped.
    fn len(&self) -> usize;
}

struct FairGroup<T, Q> {
    queue: Q,
    last_access_secs: AtomicU64,
    _phantom: PhantomData<fn() -> T>,
}

impl<T, Q: GroupInnerQueue<T>> FairGroup<T, Q> {
    pub fn new(queue: Q) -> Self {
        Self {
            queue,
            last_access_secs: AtomicU64::new(now_coarse_seconds()),
            _phantom: PhantomData,
        }
    }

    #[inline]
    pub fn push(&self, t: T) {
        self.update_last_access_secs();
        self.queue.push(t)
    }

    #[inline]
    pub fn pop(&self) -> Option<T> {
        self.update_last_access_secs();
        self.queue.pop()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.queue.len() == 0
    }

    #[inline]
    fn secs_since_last_access(&self) -> u64 {
        now_coarse_seconds().saturating_sub(self.last_access_secs.load(Ordering::Relaxed))
    }

    #[inline]
    fn update_last_access_secs(&self) {
        self.last_access_secs
            .fetch_max(now_coarse_seconds(), Ordering::Relaxed);
    }
}

/// A small concurrent registry for virtual-time based fair selection.
///
/// Each entry is identified by a key and owns a value plus an atomic virtual
/// time. Callers usually use this type by:
///
/// * creating or accessing a value with [`FairRegistry::ensure_entry_with_op`],
/// * selecting the eligible entry with the smallest virtual time through
///   [`FairRegistry::operate_min_vt_entry`],
/// * charging completed or reserved work with
///   [`FairRegistry::increase_entry_vt`], and
/// * removing idle entries with [`FairRegistry::remove_if`].
///
/// New entries start at `max_vt - max_vt_lag`, and existing entries are moved
/// forward to at least that lower bound before operations. This lets a member
/// catch up without being placed arbitrarily far behind the current maximum
/// virtual time. The operation closures run while the registry lock is held, so
/// they should be short, non-blocking, and must not call back into the same
/// registry.
///
/// A typical usage pattern is to keep one queue per member, pick the non-empty
/// queue with the smallest virtual time, and charge the selected member after
/// the work is reserved or completed:
///
/// ```ignore
/// use std::{collections::VecDeque, sync::Mutex};
///
/// use tikv_util::fairness::FairRegistry;
///
/// let registry = FairRegistry::new(1_000_000);
///
/// registry.ensure_entry_with_op(
///     1_u32,
///     || Mutex::new(VecDeque::new()),
///     |queue| queue.lock().unwrap().push_back("task-a"),
/// );
/// registry.ensure_entry_with_op(
///     2_u32,
///     || Mutex::new(VecDeque::new()),
///     |queue| queue.lock().unwrap().push_back("task-b"),
/// );
///
/// if let Some((key, task)) = registry.operate_min_vt_entry(
///     |_, queue| !queue.lock().unwrap().is_empty(),
///     |key, queue| (key, queue.lock().unwrap().pop_front().unwrap()),
/// ) {
///     // Charge either an estimated reservation or measured execution time.
///     registry.increase_entry_vt(key, 100);
///     handle_task(task);
/// }
///
/// registry.remove_if(1_u32, |queue| queue.lock().unwrap().is_empty());
/// ```
///
/// Performance characteristics:
///
/// * When the member set is stable, lookups and selections take a shared read
///   lock. Virtual-time charging updates only the matched entry and the global
///   maximum with atomics, so it does not need the registry write lock.
/// * Adding or removing members takes the write lock and blocks concurrent
///   lookups/selections while the membership change is in progress. This is
///   expected to be rare.
/// * Selection scans all registered members instead of maintaining a heap. It
///   is cheap for a small member set, but it is not suitable for very large
///   member sets or high-churn membership.
///
/// This type fits workloads with a small or moderate number of long-lived
/// members, frequent per-member operations, and infrequent joins/removals. For
/// example, it is suitable for choosing among active keyspace queues, but not
/// for scheduling millions of short-lived tasks directly.
pub struct FairRegistry<K, V> {
    /// Registered entries.
    entries: RwLock<HashMap<K, FairRegistryEntry<K, V>>>,
    /// Highest virtual time observed across all entries.
    max_vt: AtomicU64,
    /// Maximum distance a new or lagging entry may stay behind `max_vt`.
    max_vt_lag: u64,
}

impl<K, V> FairRegistry<K, V>
where
    K: FairRegistryKey,
    V: FairRegistryValue,
{
    /// Creates an empty registry.
    ///
    /// `max_vt_lag` limits how far a new entry may start behind the current
    /// maximum virtual time. Its unit is defined by the caller's virtual-time
    /// accounting.
    #[inline]
    pub fn new(max_vt_lag: u64) -> Self {
        FairRegistry {
            max_vt_lag,
            entries: RwLock::new(HashMap::new()),
            max_vt: AtomicU64::new(0),
        }
    }

    /// Returns whether `key` is currently registered.
    #[inline]
    pub fn contains_key(&self, key: K) -> bool {
        self.entries.read().unwrap().contains_key(&key)
    }

    /// Returns a snapshot of all registered keys.
    ///
    /// The returned list may become stale immediately if another thread adds or
    /// removes entries.
    #[inline]
    pub fn keys(&self) -> Vec<K> {
        self.entries.read().unwrap().keys().cloned().collect()
    }

    /// Returns a snapshot of the current entry count.
    ///
    /// The returned value may become stale immediately if another thread adds
    /// or removes entries.
    #[inline]
    pub fn entry_count(&self) -> usize {
        self.entries.read().unwrap().len()
    }

    /// Ensures an entry exists and then operates on its value.
    ///
    /// If `key` is missing, `create` builds the value and the new entry starts
    /// at the registry's current fair initial virtual time. Existing entries
    /// are first clamped to the current fair lower bound. `operate` runs while
    /// the registry lock is held, so it should be short and must not call back
    /// into this registry. Existing entries are operated under the read lock.
    /// Missing entries take the write lock, re-check the key, create the entry
    /// if still missing, and then run `operate` under that write lock. The
    /// write-lock path blocks other registry methods, including
    /// [`FairRegistry::operate_entry`], [`FairRegistry::operate_min_vt_entry`],
    /// [`FairRegistry::increase_entry_vt`], and
    /// [`FairRegistry::decrease_entry_vt`].
    pub fn ensure_entry_with_op<R>(
        &self,
        key: K,
        create: impl FnOnce() -> V,
        operate: impl FnOnce(&V) -> R,
    ) -> R {
        if let Some(entry) = self.entries.read().unwrap().get(&key) {
            self.ensure_entry_min_vt(entry);
            return operate(&entry.val);
        }

        let mut entries = self.entries.write().unwrap();
        if let Some(entry) = entries.get(&key) {
            self.ensure_entry_min_vt(entry);
            return operate(&entry.val);
        }

        let entry = entries
            .entry(key)
            .or_insert_with(|| self.new_entry(key, create()));
        self.ensure_entry_min_vt(entry);
        operate(&entry.val)
    }

    /// Operates on an existing entry's value.
    ///
    /// Returns `None` if `key` is not registered. The entry is first clamped to
    /// the current fair lower bound. `operate` runs while the registry read
    /// lock is held.
    pub fn operate_entry<R>(&self, key: K, operate: impl FnOnce(&V) -> R) -> Option<R> {
        self.entries.read().unwrap().get(&key).map(|entry| {
            self.ensure_entry_min_vt(entry);
            operate(&entry.val)
        })
    }

    /// Finds the eligible entry with the smallest virtual time and operates on
    /// it.
    ///
    /// `cond` decides whether an entry is eligible. Selection compares raw
    /// virtual time; this is equivalent to comparing clamped virtual time
    /// because the clamp is `max(vt, lower_bound)`, which preserves
    /// ordering except for ties below the lower bound. Only the selected
    /// entry is clamped before `operate` runs, avoiding writes to every
    /// lagging candidate during scans. Returns `None` when no entry
    /// satisfies `cond`. Both `cond` and `operate` run while the registry
    /// read lock is held.
    pub fn operate_min_vt_entry<R>(
        &self,
        cond: impl Fn(K, &V) -> bool,
        operate: impl FnOnce(K, &V) -> R,
    ) -> Option<R> {
        let guard = self.entries.read().unwrap();
        Self::find_entry_with_min_vt(&guard, &cond).map(|entry| {
            self.ensure_entry_min_vt(entry);
            operate(entry.key, &entry.val)
        })
    }

    /// Increases an entry's virtual time.
    ///
    /// The entry is first moved forward to at least `max_vt - max_vt_lag`, then
    /// `inc` is charged on top of that clamped value. Returns `false` if `key`
    /// is not registered.
    pub fn increase_entry_vt(&self, key: K, inc: u64) -> bool {
        let guard = self.entries.read().unwrap();
        match guard.get(&key) {
            Some(entry) => {
                self.increase_entry_ref_vt(entry, inc);
                true
            }
            None => false,
        }
    }

    /// Decreases an entry's virtual time.
    ///
    /// This is intended to correct an earlier precharge. The entry may move
    /// backward, but the registry's `max_vt` is not reduced. Returns `false` if
    /// `key` is not registered.
    pub fn decrease_entry_vt(&self, key: K, dec: u64) -> bool {
        let guard = self.entries.read().unwrap();
        match guard.get(&key) {
            Some(entry) => {
                self.decrease_entry_ref_vt(entry, dec);
                true
            }
            None => false,
        }
    }

    /// Removes an entry if it exists and satisfies `cond`.
    ///
    /// The condition is checked while holding the registry write lock. Returns
    /// the removed value, or `None` if the entry is missing or `cond` returns
    /// `false`. The write lock blocks other registry methods while `cond` runs,
    /// including [`FairRegistry::operate_entry`],
    /// [`FairRegistry::operate_min_vt_entry`],
    /// [`FairRegistry::increase_entry_vt`], and
    /// [`FairRegistry::decrease_entry_vt`].
    pub fn remove_if(&self, key: K, cond: impl Fn(&V) -> bool) -> Option<V> {
        let mut guard = self.entries.write().unwrap();
        if let Some(entry) = guard.get(&key) {
            if cond(&entry.val) {
                return guard.remove(&key).map(|entry| entry.val);
            }
        }
        None
    }

    #[inline]
    fn new_entry(&self, key: K, val: V) -> FairRegistryEntry<K, V> {
        FairRegistryEntry::new(
            key,
            val,
            self.max_vt
                .load(Ordering::Acquire)
                .saturating_sub(self.max_vt_lag),
        )
    }

    #[inline]
    fn find_entry_with_min_vt(
        entries: &HashMap<K, FairRegistryEntry<K, V>>,
        cond: impl Fn(K, &V) -> bool,
    ) -> Option<&FairRegistryEntry<K, V>> {
        entries
            .values()
            .filter(|s| cond(s.key, &s.val))
            .min_by_key(|e| e.virtual_time())
    }

    #[inline]
    fn ensure_entry_min_vt(&self, entry: &FairRegistryEntry<K, V>) {
        let old_vt = entry.virtual_time.load(Ordering::Relaxed);
        let min_vt = self
            .max_vt
            .load(Ordering::Relaxed)
            .saturating_sub(self.max_vt_lag);
        if old_vt < min_vt {
            entry.virtual_time.fetch_max(min_vt, Ordering::Relaxed);
        }
    }

    fn increase_entry_ref_vt(&self, entry: &FairRegistryEntry<K, V>, inc: u64) {
        loop {
            let old_vt = entry.virtual_time.load(Ordering::Acquire);
            let max_vt = self.max_vt.load(Ordering::Acquire);
            let min_vt = max_vt.saturating_sub(self.max_vt_lag);
            let new_vt = old_vt.max(min_vt).saturating_add(inc);
            if entry
                .virtual_time
                .compare_exchange(old_vt, new_vt, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                if new_vt > max_vt {
                    self.max_vt.fetch_max(new_vt, Ordering::AcqRel);
                }
                return;
            }
        }
    }

    fn decrease_entry_ref_vt(&self, entry: &FairRegistryEntry<K, V>, dec: u64) {
        loop {
            let old_vt = entry.virtual_time.load(Ordering::Acquire);
            let max_vt = self.max_vt.load(Ordering::Acquire);
            // Correction is bounded by the same lag window as normal fair
            // selection. Concurrent precharges may temporarily advance `max_vt`
            // by about `inflight_precharges * precharge`; callers should keep
            // each precharge much smaller than `max_vt_lag` because correction
            // can lower this entry but does not roll back the global watermark.
            let min_vt = max_vt.saturating_sub(self.max_vt_lag);
            let new_vt = old_vt.saturating_sub(dec).max(min_vt);
            if entry
                .virtual_time
                .compare_exchange(old_vt, new_vt, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                return;
            }
        }
    }
}

/// Records virtual-time charges for fair scheduling.
///
/// This trait lets task execution code update fairness accounting without
/// depending on the concrete queue or registry type.
pub trait FairnessRecorder<K: FairRegistryKey>: Send + Sync {
    /// Increases the virtual time of `k` by `vt`.
    ///
    /// The unit of `vt` is defined by the caller's fairness accounting, for
    /// example estimated reserved work or measured execution time.
    fn increase_vt(&self, k: K, vt: u64);

    /// Decreases the virtual time of `k` by `vt`.
    fn decrease_vt(&self, k: K, vt: u64);
}

impl<K: FairRegistryKey, V: FairRegistryValue> FairnessRecorder<K> for FairRegistry<K, V> {
    #[inline]
    fn increase_vt(&self, k: K, vt: u64) {
        self.increase_entry_vt(k, vt);
    }

    #[inline]
    fn decrease_vt(&self, k: K, vt: u64) {
        self.decrease_entry_vt(k, vt);
    }
}

/// Key type used by [`FairRegistry`].
pub trait FairRegistryKey: Eq + Hash + Copy + Send + Sync {}

impl<T> FairRegistryKey for T where T: Eq + Hash + Copy + Send + Sync {}

/// Value type stored in [`FairRegistry`].
pub trait FairRegistryValue: Send + Sync {}

impl<T> FairRegistryValue for T where T: Send + Sync {}

/// Internal entry stored by [`FairRegistry`].
struct FairRegistryEntry<K, V> {
    key: K,
    val: V,
    /// Virtual time used for fair selection.
    virtual_time: AtomicU64,
}

impl<K, V> FairRegistryEntry<K, V> {
    #[inline]
    fn new(key: K, val: V, initial_vt: u64) -> Self {
        Self {
            key,
            val,
            virtual_time: AtomicU64::new(initial_vt),
        }
    }

    fn virtual_time(&self) -> u64 {
        self.virtual_time.load(Ordering::Acquire)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        collections::VecDeque,
        hash::{Hash, Hasher},
        sync::{
            Arc, Barrier, Mutex,
            atomic::{AtomicBool, AtomicUsize},
            mpsc,
        },
        thread,
    };

    use super::*;

    fn entry_vt<V: FairRegistryValue>(registry: &FairRegistry<u32, V>, key: u32) -> u64 {
        registry
            .entries
            .read()
            .unwrap()
            .get(&key)
            .unwrap()
            .virtual_time()
    }

    struct TestInnerQueue<T> {
        queue: Mutex<VecDeque<T>>,
    }

    impl<T> Default for TestInnerQueue<T> {
        fn default() -> Self {
            Self {
                queue: Mutex::new(VecDeque::new()),
            }
        }
    }

    impl<T> GroupInnerQueue<T> for TestInnerQueue<T> {
        fn push(&self, t: T) {
            self.queue.lock().unwrap().push_back(t);
        }

        fn pop(&self) -> Option<T> {
            self.queue.lock().unwrap().pop_front()
        }

        fn len(&self) -> usize {
            self.queue.lock().unwrap().len()
        }
    }

    struct MaybeFalseReadyInnerQueue<T> {
        queue: Mutex<VecDeque<T>>,
        return_none_once: bool,
        pop_attempts: AtomicUsize,
    }

    impl<T> MaybeFalseReadyInnerQueue<T> {
        fn new(return_none_once: bool) -> Self {
            Self {
                queue: Mutex::new(VecDeque::new()),
                return_none_once,
                pop_attempts: AtomicUsize::new(0),
            }
        }
    }

    impl<T> GroupInnerQueue<T> for MaybeFalseReadyInnerQueue<T> {
        fn push(&self, t: T) {
            self.queue.lock().unwrap().push_back(t);
        }

        fn pop(&self) -> Option<T> {
            if self.return_none_once && self.pop_attempts.fetch_add(1, Ordering::SeqCst) == 0 {
                return None;
            }
            self.queue.lock().unwrap().pop_front()
        }

        fn len(&self) -> usize {
            self.queue.lock().unwrap().len()
        }
    }

    struct RetryPanicFalseReadyInnerQueue<T> {
        queue: Mutex<VecDeque<T>>,
        pop_attempts: AtomicUsize,
    }

    impl<T> Default for RetryPanicFalseReadyInnerQueue<T> {
        fn default() -> Self {
            Self {
                queue: Mutex::new(VecDeque::new()),
                pop_attempts: AtomicUsize::new(0),
            }
        }
    }

    impl<T> GroupInnerQueue<T> for RetryPanicFalseReadyInnerQueue<T> {
        fn push(&self, t: T) {
            self.queue.lock().unwrap().push_back(t);
        }

        fn pop(&self) -> Option<T> {
            let attempts = self.pop_attempts.fetch_add(1, Ordering::SeqCst);
            assert_eq!(attempts, 0, "pop retried a false-ready group");
            None
        }

        fn len(&self) -> usize {
            self.queue.lock().unwrap().len()
        }
    }

    fn new_test_fair_queue<T: Send + 'static>() -> MultiGroupFairQueue<u32, T, TestInnerQueue<T>> {
        MultiGroupFairQueue::new(100, |_| TestInnerQueue::default())
    }

    fn insert_empty_group<T: Send + 'static>(
        queue: &MultiGroupFairQueue<u32, T, TestInnerQueue<T>>,
        key: u32,
    ) {
        queue.registry.ensure_entry_with_op(
            key,
            || FairGroup::new(TestInnerQueue::default()),
            |_| (),
        );
    }

    fn set_group_last_access<T: Send + 'static>(
        queue: &MultiGroupFairQueue<u32, T, TestInnerQueue<T>>,
        key: u32,
        secs: u64,
    ) {
        queue
            .registry
            .operate_entry(key, |group| {
                group.last_access_secs.store(secs, Ordering::Relaxed);
            })
            .unwrap();
    }

    struct MockNowGuard;

    impl Drop for MockNowGuard {
        fn drop(&mut self) {
            set_now_for_test(None);
        }
    }

    fn mock_now_coarse_seconds(now_secs: u64) -> MockNowGuard {
        set_now_for_test(Some(now_secs));
        MockNowGuard
    }

    fn assert_contains_each_task_once(mut tasks: Vec<usize>, total: usize) {
        tasks.sort_unstable();
        assert_eq!(tasks, (0..total).collect::<Vec<_>>());
    }

    // Coordinates all workers so they observe the same first read-miss in
    // `ensure_entry_with_op`. The test key blocks during hashing, which happens
    // while `HashMap::get` is still running under the registry read lock.
    struct RecheckCoordinator {
        barrier: Barrier,
        // Only the first round of hashing should wait. Later hash calls happen in
        // the write-lock recheck, insertion, or verification path; blocking those
        // calls would wait for threads that will never arrive.
        remaining_waits: AtomicUsize,
    }

    #[derive(Clone, Copy)]
    struct RecheckKey {
        id: u32,
        coordinator: &'static RecheckCoordinator,
    }

    impl PartialEq for RecheckKey {
        fn eq(&self, other: &Self) -> bool {
            self.id == other.id
        }
    }

    impl Eq for RecheckKey {}

    impl Hash for RecheckKey {
        fn hash<H: Hasher>(&self, state: &mut H) {
            self.id.hash(state);
            // Key 1 is the contended missing key. By waiting here, all worker
            // threads stay inside the initial read-lock lookup until everyone has
            // reached the same point. After the barrier releases, they all finish
            // that lookup as a miss and then race for the write lock. The first
            // thread creates the entry; the others must hit the write-lock
            // double-check path.
            if self.id == 1
                && self
                    .coordinator
                    .remaining_waits
                    .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
                        remaining.checked_sub(1)
                    })
                    .is_ok()
            {
                self.coordinator.barrier.wait();
            }
        }
    }

    #[test]
    fn test_fair_registry_new_entry_starts_at_fair_initial_vt() {
        let registry = FairRegistry::new(100);

        // Move the registry max vt forward so a later entry has a non-zero fair
        // initial vt.
        registry.ensure_entry_with_op(1, || (), |_| ());
        registry.increase_entry_vt(1, 1_000);

        // A new entry should not start arbitrarily far behind the current max vt.
        registry.ensure_entry_with_op(2, || (), |_| ());

        assert_eq!(entry_vt(&registry, 2), 900);
    }

    #[test]
    fn test_fair_registry_ensure_entry_clamps_existing_entry_without_recreate() {
        let registry = FairRegistry::new(100);
        let create_count = AtomicUsize::new(0);

        // Create an entry that will later become stale.
        registry.ensure_entry_with_op(
            1,
            || {
                create_count.fetch_add(1, Ordering::SeqCst);
                10
            },
            |value| assert_eq!(*value, 10),
        );

        // Move max vt forward so key 1 falls behind the allowed lag.
        registry.ensure_entry_with_op(2, || 20, |_| ());
        registry.increase_entry_vt(2, 1_000);

        // Re-entering an existing entry should clamp it and must not recreate it.
        let recreate = || create_count.fetch_add(1, Ordering::SeqCst) + 99;
        let value = registry.ensure_entry_with_op(1, recreate, |value| *value);

        assert_eq!(value, 10);
        assert_eq!(create_count.load(Ordering::SeqCst), 1);
        assert_eq!(entry_vt(&registry, 1), 900);
    }

    #[test]
    fn test_fair_registry_ensure_entry_write_lock_recheck_uses_existing_entry() {
        const THREAD_COUNT: usize = 16;

        let coordinator = Box::leak(Box::new(RecheckCoordinator {
            barrier: Barrier::new(THREAD_COUNT),
            remaining_waits: AtomicUsize::new(THREAD_COUNT),
        }));
        let existing_key = RecheckKey { id: 0, coordinator };
        let contended_key = RecheckKey { id: 1, coordinator };
        let registry = Arc::new(FairRegistry::new(100));
        let create_count = Arc::new(AtomicUsize::new(0));

        // Keep the map non-empty so HashMap must hash the contended missing key.
        registry.ensure_entry_with_op(existing_key, || 0, |_| ());

        // The contended key waits in Hash while the initial read lookup is still
        // running. All workers therefore observe the same read miss before any of
        // them can acquire the write lock and create the entry.
        let handles: Vec<_> = (0..THREAD_COUNT)
            .map(|_| {
                let registry = registry.clone();
                let create_count = create_count.clone();
                thread::spawn(move || {
                    let value = registry.ensure_entry_with_op(
                        contended_key,
                        || {
                            create_count.fetch_add(1, Ordering::SeqCst);
                            42
                        },
                        |value| *value,
                    );
                    assert_eq!(value, 42);
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
        assert_eq!(create_count.load(Ordering::SeqCst), 1);
        assert_eq!(registry.entry_count(), 2);
        assert_eq!(
            registry.operate_entry(contended_key, |value| *value),
            Some(42)
        );
    }

    #[test]
    fn test_fair_registry_operate_entry_clamps_existing_entry_and_reports_missing() {
        let registry = FairRegistry::new(100);

        // Create key 1 before max vt moves so it becomes a lagging entry.
        registry.ensure_entry_with_op(1, || 10, |_| ());
        registry.ensure_entry_with_op(2, || 20, |_| ());
        registry.increase_entry_vt(2, 1_000);

        // Missing keys should stay observable as None, while existing keys are
        // clamped before operation.
        let missing: Option<u32> = registry.operate_entry(3, |value| *value);
        let value = registry.operate_entry(1, |value| *value);

        assert_eq!(missing, None);
        assert_eq!(value, Some(10));
        assert_eq!(entry_vt(&registry, 1), 900);
    }

    #[test]
    fn test_fair_registry_operate_min_vt_entry_respects_condition() {
        let registry = FairRegistry::new(10_000);

        // Give the three entries distinct virtual times.
        registry.ensure_entry_with_op(1, || (), |_| ());
        registry.ensure_entry_with_op(2, || (), |_| ());
        registry.ensure_entry_with_op(3, || (), |_| ());
        registry.increase_entry_vt(1, 10);
        registry.increase_entry_vt(2, 20);

        // The condition should filter out otherwise eligible entries and should
        // produce None when no entry passes it.
        let selected = registry.operate_min_vt_entry(|key, _| key != 3, |key, _| key);
        let none = registry.operate_min_vt_entry(|_, _| false, |key, _| key);

        assert_eq!(selected, Some(1));
        assert_eq!(none, None);
    }

    #[test]
    fn test_fair_registry_operate_min_vt_entry_clamps_selected_entry() {
        let registry = FairRegistry::new(100);

        // Create key 1 first so it becomes stale after key 2 advances max vt.
        registry.ensure_entry_with_op(1, || (), |_| ());
        registry.ensure_entry_with_op(2, || (), |_| ());
        registry.increase_entry_vt(2, 1_000);

        // Selecting key 1 should move it into the allowed lag window before the
        // caller operates on its value.
        let selected = registry.operate_min_vt_entry(|key, _| key == 1, |key, _| key);

        assert_eq!(selected, Some(1));
        assert_eq!(entry_vt(&registry, 1), 900);
    }

    #[test]
    fn test_fair_registry_increase_entry_vt_charges_after_clamp() {
        let registry = FairRegistry::new(100);

        // Make key 1 stale by advancing max vt through another entry.
        registry.ensure_entry_with_op(1, || (), |_| ());
        registry.ensure_entry_with_op(2, || (), |_| ());
        registry.increase_entry_vt(2, 1_000);

        // The charge should be added after clamping, not swallowed by the clamp.
        registry.increase_entry_vt(1, 5);

        assert_eq!(entry_vt(&registry, 1), 905);
    }

    #[test]
    fn test_fair_registry_decrease_entry_vt_keeps_max_vt_and_lag_bound() {
        let registry = FairRegistry::new(100);

        registry.ensure_entry_with_op(1, || (), |_| ());
        registry.increase_entry_vt(1, 1_000);
        registry.decrease_entry_vt(1, 300);

        assert_eq!(entry_vt(&registry, 1), 900);
        assert_eq!(registry.max_vt.load(Ordering::Acquire), 1_000);

        registry.decrease_entry_vt(1, 1_000);
        assert_eq!(entry_vt(&registry, 1), 900);
        assert_eq!(registry.max_vt.load(Ordering::Acquire), 1_000);
    }

    #[test]
    fn test_fair_registry_missing_key_paths_return_none_or_false() {
        let registry = FairRegistry::<u32, ()>::new(100);

        // Missing entries should not be created by accounting or cleanup helpers.
        assert!(!registry.increase_entry_vt(1, 10));
        assert!(!registry.decrease_entry_vt(1, 10));
        assert_eq!(registry.remove_if(1, |_| true), None);
        assert_eq!(registry.entry_count(), 0);
    }

    #[test]
    fn test_fair_registry_remove_if_removes_only_matching_entry() {
        let registry = FairRegistry::new(100);

        // Create two entries so the condition can keep one and remove the other.
        registry.ensure_entry_with_op(1, || "keep".to_string(), |_| ());
        registry.ensure_entry_with_op(2, || "remove".to_string(), |_| ());

        // Removal should happen only when the predicate matches under the write lock.
        let kept = registry.remove_if(1, |value| value == "remove");
        let removed = registry.remove_if(2, |value| value == "remove");
        let removed_lookup = registry.operate_entry(2, |value| value.clone());

        assert_eq!(kept, None);
        assert_eq!(removed.as_deref(), Some("remove"));
        assert!(registry.contains_key(1));
        assert!(!registry.contains_key(2));
        assert_eq!(registry.entry_count(), 1);
        assert_eq!(removed_lookup, None);
    }

    #[test]
    fn test_fair_registry_concurrent_ensure_entry_creates_once() {
        const THREAD_COUNT: usize = 8;

        let registry = Arc::new(FairRegistry::new(100));
        let create_count = Arc::new(AtomicUsize::new(0));
        let barrier = Arc::new(Barrier::new(THREAD_COUNT));

        // Start all threads together so they race through the read-miss path.
        let handles: Vec<_> = (0..THREAD_COUNT)
            .map(|_| {
                let registry = registry.clone();
                let create_count = create_count.clone();
                let barrier = barrier.clone();
                thread::spawn(move || {
                    barrier.wait();
                    registry.ensure_entry_with_op(
                        1,
                        || {
                            create_count.fetch_add(1, Ordering::SeqCst);
                            42
                        },
                        |value| assert_eq!(*value, 42),
                    );
                })
            })
            .collect();

        // Only one creator should win the write-lock double-check.
        for handle in handles {
            handle.join().unwrap();
        }
        assert_eq!(create_count.load(Ordering::SeqCst), 1);
        assert_eq!(registry.entry_count(), 1);
        assert!(registry.contains_key(1));
    }

    #[test]
    fn test_fair_registry_concurrent_increase_entry_vt_preserves_all_charges() {
        const THREAD_COUNT: usize = 8;
        const INCREASES_PER_THREAD: usize = 1_000;

        let registry = Arc::new(FairRegistry::new(0));
        let barrier = Arc::new(Barrier::new(THREAD_COUNT));

        // Use max_vt_lag = 0 so this test focuses on concurrent increments, not
        // lag-window clamping.
        registry.ensure_entry_with_op(1, || (), |_| ());

        // Race all threads on the same entry to exercise the CAS retry loop.
        let handles: Vec<_> = (0..THREAD_COUNT)
            .map(|_| {
                let registry = registry.clone();
                let barrier = barrier.clone();
                thread::spawn(move || {
                    barrier.wait();
                    for _ in 0..INCREASES_PER_THREAD {
                        registry.increase_entry_vt(1, 1);
                    }
                })
            })
            .collect();

        // Every successful charge should be reflected in the final virtual time.
        for handle in handles {
            handle.join().unwrap();
        }
        assert_eq!(
            entry_vt(registry.as_ref(), 1),
            (THREAD_COUNT * INCREASES_PER_THREAD) as u64
        );
    }

    #[test]
    fn test_multi_group_fair_queue_empty_state() {
        let queue = new_test_fair_queue::<u32>();

        // A new queue has no counted tasks and no registered groups.
        assert_eq!(queue.len(), 0);
        assert!(queue.is_empty());
        assert_eq!(queue.group_count(), 0);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_push_pop_updates_len_and_group_count() {
        let queue = new_test_fair_queue::<u32>();

        // Multiple pushes into the same group should not create extra groups.
        queue.push(1, 10);
        queue.push(1, 11);
        assert_eq!(queue.len(), 2);
        assert_eq!(queue.group_count(), 1);

        // Pushing another group should increase only the group count and total len.
        queue.push(2, 20);
        assert_eq!(queue.len(), 3);
        assert_eq!(queue.group_count(), 2);
        assert!(!queue.is_empty());

        // Popping all tasks should decrement len while keeping the group records.
        let mut popped = vec![
            queue.pop().unwrap(),
            queue.pop().unwrap(),
            queue.pop().unwrap(),
        ];
        popped.sort();
        assert_eq!(popped, vec![10, 11, 20]);
        assert_eq!(queue.len(), 0);
        assert!(queue.is_empty());
        assert_eq!(queue.group_count(), 2);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_pop_after_push_notification_never_returns_none() {
        const PRODUCER_COUNT: usize = 4;
        const TASKS_PER_PRODUCER: usize = 128;
        const GROUP_COUNT: usize = 4;
        let total_tasks = PRODUCER_COUNT * TASKS_PER_PRODUCER;

        let queue = Arc::new(new_test_fair_queue::<usize>());
        let start = Arc::new(Barrier::new(PRODUCER_COUNT));
        let (tx, rx) = mpsc::channel();

        // Producers notify only after push returns, matching the queue contract.
        let handles: Vec<_> = (0..PRODUCER_COUNT)
            .map(|producer| {
                let queue = queue.clone();
                let start = start.clone();
                let tx = tx.clone();
                thread::spawn(move || {
                    start.wait();
                    for index in 0..TASKS_PER_PRODUCER {
                        let task = producer * TASKS_PER_PRODUCER + index;
                        queue.push((task % GROUP_COUNT) as u32, task);
                        tx.send(task).unwrap();
                    }
                })
            })
            .collect();
        drop(tx);

        // Once a notification is received, at least one pushed task should still
        // be available for this single consumer.
        let mut popped = Vec::with_capacity(total_tasks);
        for _ in 0..total_tasks {
            rx.recv().unwrap();
            popped.push(queue.pop().expect("push-returned task should be visible"));
        }

        for handle in handles {
            handle.join().unwrap();
        }
        assert_contains_each_task_once(popped, total_tasks);
        assert_eq!(queue.len(), 0);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_concurrent_pushes_eventually_pop_all_tasks() {
        const PRODUCER_COUNT: usize = 8;
        const TASKS_PER_PRODUCER: usize = 128;
        const GROUP_COUNT: usize = 5;
        let total_tasks = PRODUCER_COUNT * TASKS_PER_PRODUCER;

        let queue = Arc::new(new_test_fair_queue::<usize>());
        let start = Arc::new(Barrier::new(PRODUCER_COUNT));

        // Concurrent producers share a small group set, so the test covers both
        // group creation races and existing-group push races.
        let handles: Vec<_> = (0..PRODUCER_COUNT)
            .map(|producer| {
                let queue = queue.clone();
                let start = start.clone();
                thread::spawn(move || {
                    start.wait();
                    for index in 0..TASKS_PER_PRODUCER {
                        let task = producer * TASKS_PER_PRODUCER + index;
                        queue.push((task % GROUP_COUNT) as u32, task);
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
        assert_eq!(queue.len(), total_tasks);
        assert_eq!(queue.group_count(), GROUP_COUNT);

        // After all pushes have returned, every task should eventually be popped.
        let mut popped = Vec::with_capacity(total_tasks);
        for _ in 0..total_tasks {
            popped.push(queue.pop().unwrap());
        }
        assert_contains_each_task_once(popped, total_tasks);
        assert_eq!(queue.len(), 0);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_concurrent_poppers_consume_each_task_once() {
        const CONSUMER_COUNT: usize = 8;
        const TOTAL_TASKS: usize = 1024;
        const GROUP_COUNT: usize = 7;

        let queue = Arc::new(new_test_fair_queue::<usize>());
        for task in 0..TOTAL_TASKS {
            queue.push((task % GROUP_COUNT) as u32, task);
        }
        assert_eq!(queue.len(), TOTAL_TASKS);

        let start = Arc::new(Barrier::new(CONSUMER_COUNT));
        let popped = Arc::new(Mutex::new(Vec::with_capacity(TOTAL_TASKS)));

        // All consumers start together and pop until the prefilled queue is empty.
        let handles: Vec<_> = (0..CONSUMER_COUNT)
            .map(|_| {
                let queue = queue.clone();
                let start = start.clone();
                let popped = popped.clone();
                thread::spawn(move || {
                    start.wait();
                    while let Some(task) = queue.pop() {
                        popped.lock().unwrap().push(task);
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
        let popped = Arc::try_unwrap(popped).unwrap().into_inner().unwrap();
        assert_contains_each_task_once(popped, TOTAL_TASKS);
        assert_eq!(queue.len(), 0);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_preserves_inner_queue_order_within_group() {
        let queue = new_test_fair_queue::<u32>();

        // With only one group, pop order should be the inner queue FIFO order.
        queue.push(1, 10);
        queue.push(1, 11);
        queue.push(1, 12);

        assert_eq!(queue.pop(), Some(10));
        assert_eq!(queue.pop(), Some(11));
        assert_eq!(queue.pop(), Some(12));
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_pops_group_with_smaller_vt() {
        let queue = new_test_fair_queue::<&'static str>();

        // Make group 1 more expensive so group 2 has the smaller virtual time.
        queue.push(1, "group-1");
        queue.push(2, "group-2");
        FairnessRecorder::increase_vt(&queue, 1, 10);

        assert_eq!(queue.pop(), Some("group-2"));
        assert_eq!(queue.pop(), Some("group-1"));
    }

    #[test]
    fn test_multi_group_fair_queue_skips_empty_group_even_if_its_vt_is_smaller() {
        let queue = new_test_fair_queue::<&'static str>();

        // Leave group 1 registered but empty with a smaller virtual time.
        queue.push(1, "empty-first");
        assert_eq!(queue.pop(), Some("empty-first"));

        // Group 2 has a task and a larger virtual time, so pop must ignore group 1.
        queue.push(2, "ready");
        FairnessRecorder::increase_vt(&queue, 2, 10);

        assert_eq!(queue.pop(), Some("ready"));
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_drain_removes_tasks_but_keeps_groups() {
        let queue = new_test_fair_queue::<u32>();

        // Put tasks into multiple groups so drain has to visit the key snapshot.
        queue.push(1, 10);
        queue.push(1, 11);
        queue.push(2, 20);
        assert_eq!(queue.len(), 3);
        assert_eq!(queue.group_count(), 2);

        // Drain removes all current tasks but does not delete the group entries.
        queue.drain();

        assert_eq!(queue.len(), 0);
        assert!(queue.is_empty());
        assert_eq!(queue.group_count(), 2);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_remove_inactive_groups_removes_only_empty_expired_groups() {
        let queue = new_test_fair_queue::<u32>();
        let mock_now_secs = 100;
        let old_secs = 90;
        let _mock_now = mock_now_coarse_seconds(mock_now_secs);

        // Build one empty expired group, one non-empty expired group, and one
        // empty recent group with a mocked current time.
        insert_empty_group(&queue, 1);
        queue.push(2, 20);
        insert_empty_group(&queue, 3);
        set_group_last_access(&queue, 1, old_secs);
        set_group_last_access(&queue, 2, old_secs);
        set_group_last_access(&queue, 3, mock_now_secs);

        // Cleanup should remove only the group that is both empty and expired.
        queue.remove_inactive_groups(5);

        assert!(!queue.registry.contains_key(1));
        assert!(queue.registry.contains_key(2));
        assert!(queue.registry.contains_key(3));
        assert_eq!(queue.group_count(), 2);
        assert_eq!(queue.len(), 1);
        assert_eq!(queue.pop(), Some(20));
    }

    #[test]
    fn test_multi_group_fair_queue_cleanup_with_concurrent_push_pop_does_not_lose_tasks() {
        const PRODUCER_COUNT: usize = 4;
        const TASKS_PER_PRODUCER: usize = 128;
        const GROUP_COUNT: usize = 8;
        let total_tasks = PRODUCER_COUNT * TASKS_PER_PRODUCER;

        let queue = Arc::new(new_test_fair_queue::<usize>());
        let start = Arc::new(Barrier::new(PRODUCER_COUNT));
        let cleanup_done = Arc::new(AtomicBool::new(false));
        let (tx, rx) = mpsc::channel();

        // Cleanup sees a later mocked time, so groups that become empty are
        // eligible for removal while producers and the consumer keep using them.
        let cleanup_handle = {
            let queue = queue.clone();
            let cleanup_done = cleanup_done.clone();
            thread::spawn(move || {
                let _mock_now = mock_now_coarse_seconds(200);
                while !cleanup_done.load(Ordering::Acquire) {
                    queue.remove_inactive_groups(0);
                    thread::yield_now();
                }
                queue.remove_inactive_groups(0);
            })
        };

        let producer_handles: Vec<_> = (0..PRODUCER_COUNT)
            .map(|producer| {
                let queue = queue.clone();
                let start = start.clone();
                let tx = tx.clone();
                thread::spawn(move || {
                    let _mock_now = mock_now_coarse_seconds(100);
                    start.wait();
                    for index in 0..TASKS_PER_PRODUCER {
                        let task = producer * TASKS_PER_PRODUCER + index;
                        queue.push((task % GROUP_COUNT) as u32, task);
                        tx.send(task).unwrap();
                        thread::yield_now();
                    }
                })
            })
            .collect();
        drop(tx);

        // Pop only after push-return notifications. This keeps the test aligned
        // with the queue's bounded consistency contract while cleanup races with
        // both producers and the consumer.
        let _mock_now = mock_now_coarse_seconds(100);
        let mut popped = Vec::with_capacity(total_tasks);
        for _ in 0..total_tasks {
            rx.recv().unwrap();
            popped.push(queue.pop().expect("cleanup should not lose pushed tasks"));
            thread::yield_now();
        }

        for handle in producer_handles {
            handle.join().unwrap();
        }
        cleanup_done.store(true, Ordering::Release);
        cleanup_handle.join().unwrap();

        assert_contains_each_task_once(popped, total_tasks);
        assert_eq!(queue.len(), 0);
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_multi_group_fair_queue_skips_false_ready_group_during_pop() {
        let create_count = Arc::new(AtomicUsize::new(0));
        let queue = {
            let create_count = create_count.clone();
            MultiGroupFairQueue::new(100, move |_| {
                let index = create_count.fetch_add(1, Ordering::SeqCst);
                MaybeFalseReadyInnerQueue::new(index == 0)
            })
        };

        // The first created group reports non-empty but returns None on its first pop.
        queue.push(1, "false-ready");
        queue.push(2, "ready");
        FairnessRecorder::increase_vt(&queue, 2, 10);

        // Pop should skip the false-ready group once and return work from group 2.
        assert_eq!(queue.pop(), Some("ready"));
        assert_eq!(queue.len(), 1);

        // The skipped group still owns its task and can be popped later.
        assert_eq!(queue.pop(), Some("false-ready"));
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn test_multi_group_fair_queue_pop_does_not_retry_skipped_false_ready_group() {
        let queue = MultiGroupFairQueue::new(100, |_| RetryPanicFalseReadyInnerQueue::default());

        // The inner queue keeps reporting non-empty but returns None. The outer
        // pop must skip that group after one failed attempt; retrying it in the
        // same call would either spin forever or hit the assertion in pop().
        queue.push(1, "not-ready");
        assert_eq!(queue.pop(), None);
        assert_eq!(
            queue
                .registry
                .operate_entry(1, |group| group.queue.pop_attempts.load(Ordering::SeqCst)),
            Some(1)
        );
        assert_eq!(queue.len(), 1);
    }

    #[test]
    fn test_multi_group_fair_queue_returns_none_when_all_candidates_are_false_ready() {
        let queue = MultiGroupFairQueue::new(100, |_| MaybeFalseReadyInnerQueue::new(true));

        // This inner queue intentionally reports non-empty but returns None once.
        // With no other ready group, pop should skip it and then report no
        // currently available task without changing the counted length.
        queue.push(1, "delayed");
        assert_eq!(queue.pop(), None);
        assert_eq!(queue.len(), 1);

        // The skipped task is still owned by the group and can be popped later.
        assert_eq!(queue.pop(), Some("delayed"));
        assert_eq!(queue.len(), 0);
    }
}