// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    collections::VecDeque,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ScanPriority {
    Low,
    High,
}

impl ScanPriority {
    pub(crate) fn as_str(self) -> &'static str {
        match self {
            ScanPriority::Low => "low",
            ScanPriority::High => "high",
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ScanSlotKind {
    Normal,
    EmergencyHigh,
}

impl ScanSlotKind {
    pub(crate) fn as_str(self) -> &'static str {
        match self {
            ScanSlotKind::Normal => "normal",
            ScanSlotKind::EmergencyHigh => "emergency_high",
        }
    }
}

struct QueuedScan<T> {
    enqueue_at: Instant,
    payload: T,
    admission: ScanAdmissionPermit,
}

impl<T> QueuedScan<T> {
    fn new(payload: T, admission: ScanAdmissionPermit) -> Self {
        Self {
            enqueue_at: Instant::now(),
            payload,
            admission,
        }
    }
}

/// RAII ownership of admission capacity for one queued or running scan.
///
/// Admission is released by the worker before the endpoint consumes slot
/// completion. Keeping it separate from `ScanPermit` preserves rejection
/// behavior when the endpoint task queue is backlogged.
pub(crate) struct ScanAdmissionPermit {
    admitted: Arc<AtomicUsize>,
}

impl Drop for ScanAdmissionPermit {
    fn drop(&mut self) {
        // Admission does not publish task data. The endpoint serializes increments,
        // and each unique permit contributes at most one decrement, so a stale
        // relaxed read may delay admission but cannot let it exceed the limit.
        self.admitted
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
                count.checked_sub(1)
            })
            .expect("scan admission released without an admitted task");
    }
}

/// Proof that a dispatched scan owns one scheduler slot.
///
/// The permit is intentionally not cloneable, so a slot can be completed at
/// most once through the scheduler API.
#[must_use = "dropping a scan permit without completing it leaks a scheduler slot"]
pub(crate) struct ScanPermit {
    slot_kind: ScanSlotKind,
}

impl ScanPermit {
    pub(crate) fn slot_kind(&self) -> &'static str {
        self.slot_kind.as_str()
    }
}

#[must_use = "dispatched scan tasks must be run or completed"]
pub(crate) struct DispatchedScanTask<T> {
    pub(crate) payload: T,
    pub(crate) priority: ScanPriority,
    pub(crate) queue_time: Duration,
    pub(crate) permit: ScanPermit,
    pub(crate) admission: ScanAdmissionPermit,
}

/// Runnable scans and the number of canceled queued scans removed while
/// finding them.
#[must_use = "dispatch results must be run and cancellation metrics recorded"]
pub(crate) struct DispatchBatch<T> {
    pub(crate) tasks: Vec<DispatchedScanTask<T>>,
    pub(crate) cancelled: usize,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct ScanSchedulerStats {
    pub(crate) low_queued: usize,
    pub(crate) high_queued: usize,
    pub(crate) normal_running: usize,
    pub(crate) emergency_high_running: usize,
    pub(crate) normal_slots: usize,
    pub(crate) emergency_high_slots: usize,
    pub(crate) admission_limit: usize,
    pub(crate) admitted: usize,
}

impl ScanSchedulerStats {
    pub(crate) fn admitted(self) -> usize {
        self.admitted
    }
}

pub(crate) struct ScanScheduler<T> {
    high_queue: VecDeque<QueuedScan<T>>,
    low_queue: VecDeque<QueuedScan<T>>,
    normal_running: usize,
    emergency_high_running: usize,
    normal_slots: usize,
    emergency_high_slots: usize,
    admission_limit: usize,
    admitted: Arc<AtomicUsize>,
    is_cancelled: fn(&T) -> bool,
}

impl<T> ScanScheduler<T> {
    pub(crate) fn new(
        normal_slots: usize,
        emergency_high_slots: usize,
        admission_limit: usize,
        is_cancelled: fn(&T) -> bool,
    ) -> Self {
        Self {
            high_queue: VecDeque::new(),
            low_queue: VecDeque::new(),
            normal_running: 0,
            emergency_high_running: 0,
            normal_slots,
            emergency_high_slots,
            admission_limit,
            admitted: Arc::default(),
            is_cancelled,
        }
    }

    pub(crate) fn can_admit(&self, priority: ScanPriority) -> bool {
        self.admitted() < self.admission_limit(priority)
    }

    pub(crate) fn admission_limit(&self, priority: ScanPriority) -> usize {
        self.admission_limit.saturating_add(match priority {
            ScanPriority::Low => 0,
            ScanPriority::High => self.emergency_high_slots,
        })
    }

    /// Acquires admission capacity, enqueues the scan, and returns all newly
    /// runnable tasks.
    ///
    /// The endpoint calls `can_admit` before it mutates registration state. The
    /// endpoint event loop is single-threaded, so scheduler state cannot change
    /// between that check and this enqueue operation.
    pub(crate) fn enqueue(&mut self, priority: ScanPriority, payload: T) -> DispatchBatch<T> {
        assert!(
            self.can_admit(priority),
            "scan task must pass admission before enqueue"
        );
        let admission_limit = self.admission_limit(priority);
        self.admitted
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
                (count < admission_limit).then_some(count + 1)
            })
            .expect("scan task must pass admission before enqueue");
        let task = QueuedScan::new(
            payload,
            ScanAdmissionPermit {
                admitted: self.admitted.clone(),
            },
        );
        match priority {
            ScanPriority::High => self.high_queue.push_back(task),
            ScanPriority::Low => self.low_queue.push_back(task),
        }
        self.dispatch_ready()
    }

    pub(crate) fn complete(&mut self, permit: ScanPermit) -> DispatchBatch<T> {
        self.release_slot(permit);
        self.dispatch_ready()
    }

    fn release_slot(&mut self, permit: ScanPermit) {
        match permit.slot_kind {
            ScanSlotKind::Normal => {
                self.normal_running = self
                    .normal_running
                    .checked_sub(1)
                    .expect("normal scan slot completed without a running task");
            }
            ScanSlotKind::EmergencyHigh => {
                self.emergency_high_running = self
                    .emergency_high_running
                    .checked_sub(1)
                    .expect("emergency scan slot completed without a running task");
            }
        }
    }

    pub(crate) fn reconfigure(
        &mut self,
        normal_slots: usize,
        emergency_high_slots: usize,
        admission_limit: usize,
    ) -> DispatchBatch<T> {
        self.normal_slots = normal_slots;
        self.emergency_high_slots = emergency_high_slots;
        self.admission_limit = admission_limit;
        self.dispatch_ready()
    }

    /// Removes all canceled queued scans while preserving the order of the
    /// remaining scans.
    pub(crate) fn remove_cancelled_queued(&mut self) -> usize {
        let is_cancelled = self.is_cancelled;
        Self::remove_cancelled_from_queue(&mut self.high_queue, is_cancelled)
            + Self::remove_cancelled_from_queue(&mut self.low_queue, is_cancelled)
    }

    fn remove_cancelled_from_queue(
        queue: &mut VecDeque<QueuedScan<T>>,
        is_cancelled: fn(&T) -> bool,
    ) -> usize {
        let original_len = queue.len();
        queue.retain(|task| !is_cancelled(&task.payload));
        original_len - queue.len()
    }

    pub(crate) fn stats(&self) -> ScanSchedulerStats {
        ScanSchedulerStats {
            low_queued: self.low_queue.len(),
            high_queued: self.high_queue.len(),
            normal_running: self.normal_running,
            emergency_high_running: self.emergency_high_running,
            normal_slots: self.normal_slots,
            emergency_high_slots: self.emergency_high_slots,
            admission_limit: self.admission_limit,
            admitted: self.admitted(),
        }
    }

    fn admitted(&self) -> usize {
        self.admitted.load(Ordering::Relaxed)
    }

    fn total_running(&self) -> usize {
        self.normal_running
            .saturating_add(self.emergency_high_running)
    }

    fn total_slots(&self) -> usize {
        self.normal_slots.saturating_add(self.emergency_high_slots)
    }

    fn dispatch_ready(&mut self) -> DispatchBatch<T> {
        let mut tasks = Vec::new();
        let mut cancelled = 0;
        while self.normal_running < self.normal_slots && self.total_running() < self.total_slots() {
            let task = self
                .high_queue
                .pop_front()
                .map(|task| (ScanPriority::High, task))
                .or_else(|| {
                    self.low_queue
                        .pop_front()
                        .map(|task| (ScanPriority::Low, task))
                });
            let Some((priority, task)) = task else {
                break;
            };
            if (self.is_cancelled)(&task.payload) {
                cancelled += 1;
                continue;
            }
            self.normal_running += 1;
            tasks.push(Self::dispatch(task, priority, ScanSlotKind::Normal));
        }

        while self.emergency_high_running < self.emergency_high_slots {
            // After a config shrink, normal running tasks can temporarily exceed
            // the new limit. Emergency work starts only after the normal count
            // drains back to the configured saturation point.
            if self.normal_running != self.normal_slots {
                break;
            }
            let Some(task) = self.high_queue.pop_front() else {
                break;
            };
            if (self.is_cancelled)(&task.payload) {
                cancelled += 1;
                continue;
            }
            self.emergency_high_running += 1;
            tasks.push(Self::dispatch(
                task,
                ScanPriority::High,
                ScanSlotKind::EmergencyHigh,
            ));
        }
        DispatchBatch { tasks, cancelled }
    }

    fn dispatch(
        task: QueuedScan<T>,
        priority: ScanPriority,
        slot_kind: ScanSlotKind,
    ) -> DispatchedScanTask<T> {
        DispatchedScanTask {
            payload: task.payload,
            priority,
            queue_time: task.enqueue_at.elapsed(),
            permit: ScanPermit { slot_kind },
            admission: task.admission,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn scheduler(
        normal_slots: usize,
        emergency_high_slots: usize,
        admission_limit: usize,
    ) -> ScanScheduler<u64> {
        ScanScheduler::new(
            normal_slots,
            emergency_high_slots,
            admission_limit,
            |payload| *payload == 0,
        )
    }

    fn into_tasks(batch: DispatchBatch<u64>) -> Vec<DispatchedScanTask<u64>> {
        assert_eq!(batch.cancelled, 0);
        batch.tasks
    }

    fn enqueue(
        scheduler: &mut ScanScheduler<u64>,
        priority: ScanPriority,
        payload: u64,
    ) -> Vec<DispatchedScanTask<u64>> {
        into_tasks(scheduler.enqueue(priority, payload))
    }

    fn dispatched(tasks: &[DispatchedScanTask<u64>]) -> Vec<(u64, ScanPriority, &'static str)> {
        tasks
            .iter()
            .map(|task| (task.payload, task.priority, task.permit.slot_kind()))
            .collect()
    }

    fn complete(
        scheduler: &mut ScanScheduler<u64>,
        mut tasks: Vec<DispatchedScanTask<u64>>,
    ) -> Vec<DispatchedScanTask<u64>> {
        assert_eq!(tasks.len(), 1);
        into_tasks(scheduler.complete(tasks.pop().unwrap().permit))
    }

    #[test]
    fn test_low_tasks_fill_normal_slots() {
        let mut scheduler = scheduler(2, 1, 3);
        let mut tasks = enqueue(&mut scheduler, ScanPriority::Low, 1);
        tasks.extend(enqueue(&mut scheduler, ScanPriority::Low, 2));
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 3).is_empty());

        assert_eq!(
            dispatched(&tasks),
            vec![
                (1, ScanPriority::Low, "normal"),
                (2, ScanPriority::Low, "normal")
            ]
        );
        assert_eq!(
            scheduler.stats(),
            ScanSchedulerStats {
                low_queued: 1,
                high_queued: 0,
                normal_running: 2,
                emergency_high_running: 0,
                normal_slots: 2,
                emergency_high_slots: 1,
                admission_limit: 3,
                admitted: 3,
            }
        );
    }

    #[test]
    fn test_high_task_uses_normal_slot_when_available() {
        let mut scheduler = scheduler(2, 1, 3);
        let tasks = enqueue(&mut scheduler, ScanPriority::High, 1);

        assert_eq!(dispatched(&tasks), vec![(1, ScanPriority::High, "normal")]);
        assert_eq!(scheduler.stats().emergency_high_running, 0);
    }

    #[test]
    fn test_high_task_dispatches_before_queued_low_task() {
        let mut scheduler = scheduler(1, 0, 3);
        let running = enqueue(&mut scheduler, ScanPriority::Low, 1);
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 2).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 3).is_empty());

        let tasks = complete(&mut scheduler, running);

        assert_eq!(dispatched(&tasks), vec![(3, ScanPriority::High, "normal")]);
        assert_eq!(scheduler.stats().low_queued, 1);
    }

    #[test]
    fn test_cancelled_task_does_not_take_normal_slot() {
        let mut scheduler = scheduler(1, 0, 1);

        let batch = scheduler.enqueue(ScanPriority::Low, 0);

        assert!(batch.tasks.is_empty());
        assert_eq!(batch.cancelled, 1);
        assert_eq!(scheduler.stats().admitted(), 0);
        assert_eq!(scheduler.stats().normal_running, 0);
    }

    #[test]
    fn test_complete_skips_cancelled_prefix() {
        let mut scheduler = scheduler(1, 0, 4);
        let running = enqueue(&mut scheduler, ScanPriority::Low, 1);
        assert!(enqueue(&mut scheduler, ScanPriority::High, 0).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 0).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 3).is_empty());

        let batch = scheduler.complete(running.into_iter().next().unwrap().permit);

        assert_eq!(batch.cancelled, 2);
        assert_eq!(
            dispatched(&batch.tasks),
            vec![(3, ScanPriority::High, "normal")]
        );
        assert_eq!(scheduler.stats().admitted(), 1);
        assert_eq!(scheduler.stats().normal_running, 1);
    }

    #[test]
    fn test_dispatch_ready_skips_large_cancelled_prefix() {
        const CANCELLED_TASKS: usize = 10_000;
        let mut scheduler = scheduler(0, 0, CANCELLED_TASKS + 1);
        for _ in 0..CANCELLED_TASKS {
            assert!(enqueue(&mut scheduler, ScanPriority::Low, 0).is_empty());
        }
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 1).is_empty());

        let batch = scheduler.reconfigure(1, 0, CANCELLED_TASKS + 1);

        assert_eq!(batch.cancelled, CANCELLED_TASKS);
        assert_eq!(
            dispatched(&batch.tasks),
            vec![(1, ScanPriority::Low, "normal")]
        );
        assert_eq!(scheduler.stats().admitted(), 1);
        assert_eq!(scheduler.stats().normal_running, 1);
    }

    #[test]
    fn test_high_task_uses_emergency_slot_when_normal_slots_are_full() {
        let mut scheduler = scheduler(2, 1, 3);
        let mut running = enqueue(&mut scheduler, ScanPriority::Low, 1);
        running.extend(enqueue(&mut scheduler, ScanPriority::Low, 2));

        let tasks = enqueue(&mut scheduler, ScanPriority::High, 3);

        assert_eq!(
            dispatched(&tasks),
            vec![(3, ScanPriority::High, "emergency_high")]
        );
        assert_eq!(scheduler.stats().normal_running, 2);
        assert_eq!(scheduler.stats().emergency_high_running, 1);
    }

    #[test]
    fn test_emergency_slot_never_runs_low_task() {
        let mut scheduler = scheduler(1, 1, 2);
        let _running = enqueue(&mut scheduler, ScanPriority::Low, 1);

        assert!(enqueue(&mut scheduler, ScanPriority::Low, 2).is_empty());
        assert_eq!(scheduler.stats().low_queued, 1);
    }

    #[test]
    fn test_cancelled_high_task_does_not_take_emergency_slot() {
        let mut scheduler = scheduler(1, 1, 2);
        let _normal = enqueue(&mut scheduler, ScanPriority::Low, 1);

        let cancelled = scheduler.enqueue(ScanPriority::High, 0);
        assert!(cancelled.tasks.is_empty());
        assert_eq!(cancelled.cancelled, 1);
        assert_eq!(scheduler.stats().emergency_high_running, 0);

        let emergency = enqueue(&mut scheduler, ScanPriority::High, 2);
        assert_eq!(
            dispatched(&emergency),
            vec![(2, ScanPriority::High, "emergency_high")]
        );
    }

    #[test]
    fn test_cancel_queued_tasks_preserves_priority_order() {
        let mut scheduler = scheduler(0, 0, 4);
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 1).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 0).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 3).is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 0).is_empty());

        assert_eq!(scheduler.remove_cancelled_queued(), 2);
        assert_eq!(scheduler.stats().admitted(), 2);

        let tasks = into_tasks(scheduler.reconfigure(2, 0, 4));
        assert_eq!(
            dispatched(&tasks),
            vec![
                (3, ScanPriority::High, "normal"),
                (1, ScanPriority::Low, "normal")
            ]
        );
    }

    #[test]
    fn test_zero_emergency_slots_disable_active_scan_expansion() {
        let mut scheduler = scheduler(1, 0, 2);
        let running = enqueue(&mut scheduler, ScanPriority::Low, 1);
        assert!(enqueue(&mut scheduler, ScanPriority::High, 2).is_empty());

        let tasks = complete(&mut scheduler, running);

        assert_eq!(dispatched(&tasks), vec![(2, ScanPriority::High, "normal")]);
    }

    #[test]
    fn test_shrinking_limits_waits_for_running_tasks_to_complete() {
        let mut scheduler = scheduler(2, 1, 4);
        let mut normal = enqueue(&mut scheduler, ScanPriority::Low, 1);
        normal.extend(enqueue(&mut scheduler, ScanPriority::Low, 2));
        let emergency = enqueue(&mut scheduler, ScanPriority::High, 3);

        assert!(scheduler.reconfigure(1, 0, 4).tasks.is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 4).is_empty());
        assert!(complete(&mut scheduler, emergency).is_empty());
        assert!(
            scheduler
                .complete(normal.pop().unwrap().permit)
                .tasks
                .is_empty()
        );
        let tasks = into_tasks(scheduler.complete(normal.pop().unwrap().permit));

        assert_eq!(dispatched(&tasks), vec![(4, ScanPriority::High, "normal")]);
    }

    #[test]
    fn test_emergency_slot_waits_when_normal_running_exceeds_shrunk_limit() {
        let mut scheduler = scheduler(2, 0, 3);
        let mut normal = enqueue(&mut scheduler, ScanPriority::Low, 1);
        normal.extend(enqueue(&mut scheduler, ScanPriority::Low, 2));

        assert!(scheduler.reconfigure(1, 1, 3).tasks.is_empty());
        assert!(enqueue(&mut scheduler, ScanPriority::High, 3).is_empty());
        let tasks = into_tasks(scheduler.complete(normal.pop().unwrap().permit));

        assert_eq!(
            dispatched(&tasks),
            vec![(3, ScanPriority::High, "emergency_high")]
        );
    }

    #[test]
    fn test_reconfigure_does_not_exceed_new_total_slot_limit() {
        let mut scheduler = scheduler(4, 2, 8);
        let mut normal = Vec::new();
        for payload in 1..=4 {
            normal.extend(enqueue(&mut scheduler, ScanPriority::Low, payload));
        }
        let mut emergency = enqueue(&mut scheduler, ScanPriority::High, 5);
        emergency.extend(enqueue(&mut scheduler, ScanPriority::High, 6));
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 7).is_empty());

        assert!(scheduler.reconfigure(5, 0, 8).tasks.is_empty());
        assert_eq!(scheduler.total_running(), 6);

        assert!(
            scheduler
                .complete(emergency.pop().unwrap().permit)
                .tasks
                .is_empty()
        );
        assert_eq!(scheduler.total_running(), 5);

        let batch = scheduler.complete(emergency.pop().unwrap().permit);
        assert_eq!(
            dispatched(&batch.tasks),
            vec![(7, ScanPriority::Low, "normal")]
        );
        assert_eq!(scheduler.total_running(), 5);
    }

    #[test]
    fn test_total_slots_saturates() {
        let scheduler = scheduler(usize::MAX, usize::MAX, usize::MAX);

        assert_eq!(scheduler.total_slots(), usize::MAX);
    }

    #[test]
    fn test_admission_limits_include_emergency_capacity_only_for_high_priority() {
        let mut scheduler = scheduler(0, 1, 1);
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 1).is_empty());

        assert!(!scheduler.can_admit(ScanPriority::Low));
        assert!(scheduler.can_admit(ScanPriority::High));
        let _running = enqueue(&mut scheduler, ScanPriority::High, 2);
        assert!(!scheduler.can_admit(ScanPriority::High));
    }

    #[test]
    fn test_worker_completion_releases_admission_before_slot() {
        let mut scheduler = scheduler(1, 0, 1);
        let mut running = enqueue(&mut scheduler, ScanPriority::Low, 1);
        let task = running.pop().unwrap();
        let DispatchedScanTask {
            permit, admission, ..
        } = task;

        drop(admission);
        assert!(scheduler.can_admit(ScanPriority::Low));
        assert!(enqueue(&mut scheduler, ScanPriority::Low, 2).is_empty());

        let tasks = into_tasks(scheduler.complete(permit));
        assert_eq!(dispatched(&tasks), vec![(2, ScanPriority::Low, "normal")]);
    }
}