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

use std::{
    cmp::Reverse,
    collections::BinaryHeap,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc::{self, Receiver, RecvTimeoutError, Sender, TryRecvError},
    },
    thread,
    time::{Duration, Instant},
};

use anyhow::{Context, Result, ensure};
use cpu_time::ThreadTime;
use kvproto::kvrpcpb::CommandPri;
use rand::{Rng, SeedableRng, rngs::StdRng};
use tikv::read_pool::{FairCpuTaskType, ReadPoolHandle, ReadPoolTaskOptions, YatpPollEventTracker};
use tikv_util::yatp_pool::{FairCpuBudgetAction, fair_cpu_checkpoint};

use crate::{
    config::{ArrivalMode, Priority, SimulatorConfig, TaskConfig},
    metrics::{
        KeyspaceDescriptor, KeyspaceMetricStore, MetricsRecorder, ReportData, TaskDescriptor,
        WorkerCpuTimeSource,
    },
    slow_task::{TaskCompletion, TaskObservers},
};

const MAX_READY_COMPLETIONS_PER_TICK: usize = 1_024;
const CPU_TIME_CHECK_BATCH_SIZE: usize = 64;

#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
struct GeneratorEvent {
    at: Duration,
    runtime_id: usize,
}

#[derive(Debug)]
struct TaskRuntime {
    runtime_id: usize,
    /// Configured position used to keep jitter stable when another task is
    /// disabled.
    sample_identity_id: u64,
    keyspace_id: u32,
    config: Arc<TaskConfig>,
    start_time: Duration,
    end_time: Duration,
    total_arrivals: u64,
    outstanding: usize,
    next_offer_sequence: u64,
}

struct WorkloadPlan {
    keyspaces: Vec<KeyspaceDescriptor>,
    tasks: Vec<TaskDescriptor>,
    runtimes: Vec<TaskRuntime>,
}

impl TaskRuntime {
    fn arrival_deadline(&self, offer_sequence: u64) -> Duration {
        debug_assert!(offer_sequence < self.total_arrivals);
        let qps = self
            .config
            .qps
            .expect("validated open-loop task should have qps");
        self.start_time
            .checked_add(Duration::from_secs_f64(offer_sequence as f64 / qps))
            .expect("validated open-loop deadline should not overflow")
    }

    fn arrivals_due(&self, now: Duration) -> u64 {
        debug_assert_eq!(self.config.arrival_mode, ArrivalMode::OpenLoop);
        let mut first = self.next_offer_sequence;
        let mut end = self.total_arrivals;
        while first < end {
            let middle = first + (end - first) / 2;
            if self.arrival_deadline(middle) <= now {
                first = middle + 1;
            } else {
                end = middle;
            }
        }
        first
    }
}

struct WorkloadRunner<'a> {
    config: &'a SimulatorConfig,
    handle: &'a ReadPoolHandle,
    runtimes: Vec<TaskRuntime>,
    events: BinaryHeap<Reverse<GeneratorEvent>>,
    completion_tx: Sender<TaskCompletion>,
    completion_rx: Receiver<TaskCompletion>,
    deferred_completion: Option<TaskCompletion>,
    cancelled: Arc<AtomicBool>,
    next_task_id: u64,
    simulation_started_at: Option<Instant>,
    metrics: MetricsRecorder,
}

pub fn run(
    config: &SimulatorConfig,
    handle: &ReadPoolHandle,
    metric_store: KeyspaceMetricStore,
    worker_cpu_time: Arc<dyn WorkerCpuTimeSource>,
) -> Result<ReportData> {
    WorkloadRunner::new(config, handle, metric_store, worker_cpu_time)?.run()
}

impl<'a> WorkloadRunner<'a> {
    fn new(
        config: &'a SimulatorConfig,
        handle: &'a ReadPoolHandle,
        metric_store: KeyspaceMetricStore,
        worker_cpu_time: Arc<dyn WorkerCpuTimeSource>,
    ) -> Result<Self> {
        let (completion_tx, completion_rx) = mpsc::channel();
        let plan = build_workload_plan(config)?;
        let metrics = MetricsRecorder::new(
            config,
            plan.keyspaces,
            plan.tasks,
            metric_store,
            worker_cpu_time,
        )?;
        let runtimes = plan.runtimes;
        let events = runtimes
            .iter()
            .map(|runtime| {
                Reverse(GeneratorEvent {
                    at: runtime.start_time,
                    runtime_id: runtime.runtime_id,
                })
            })
            .collect();
        Ok(Self {
            config,
            handle,
            runtimes,
            events,
            completion_tx,
            completion_rx,
            deferred_completion: None,
            cancelled: Arc::new(AtomicBool::new(false)),
            next_task_id: 0,
            simulation_started_at: None,
            metrics,
        })
    }

    fn run(mut self) -> Result<ReportData> {
        let started_at = Instant::now();
        self.simulation_started_at = Some(started_at);
        let generation_end = self.config.duration.0;

        while started_at.elapsed() < generation_end {
            let mut now = started_at.elapsed();
            if now >= generation_end {
                break;
            }
            if now >= self.metrics.next_capture_deadline() {
                self.metrics.capture_if_due(now)?;
                continue;
            }
            if let Some(completion) = self.deferred_completion.take() {
                self.process_completion(completion, now, generation_end)?;
                continue;
            }

            self.process_ready_completions(started_at, generation_end)?;
            now = started_at.elapsed();
            if now >= generation_end {
                break;
            }
            if now >= self.metrics.next_capture_deadline() {
                continue;
            }
            self.process_due_events(started_at, generation_end)?;

            let now = started_at.elapsed();
            if now >= generation_end {
                break;
            }
            if now >= self.metrics.next_capture_deadline() {
                continue;
            }
            let next_event_at = self
                .events
                .peek()
                .map(|event| event.0.at)
                .unwrap_or(generation_end)
                .min(generation_end);
            let next_wakeup_at = next_event_at.min(self.metrics.next_capture_deadline());
            let wait = next_wakeup_at.saturating_sub(now);
            match self.completion_rx.recv_timeout(wait) {
                Ok(completion) => {
                    let now = started_at.elapsed();
                    if now >= generation_end || now >= self.metrics.next_capture_deadline() {
                        self.defer_completion(completion);
                    } else {
                        self.process_completion(completion, now, generation_end)?;
                    }
                }
                Err(RecvTimeoutError::Timeout) => {}
                Err(RecvTimeoutError::Disconnected) => {
                    anyhow::bail!("workload completion channel disconnected unexpectedly")
                }
            }
        }

        // Count every configured open-loop arrival even if the controller was
        // too busy to process its deadline before generation stopped.
        self.suppress_unprocessed_open_loop_arrivals();
        let generation_end_elapsed = started_at.elapsed();
        let has_drain_work = self.outstanding_task_count() > 0;
        self.metrics
            .capture_generation_end(generation_end_elapsed)?;
        if !self.config.drain_on_stop {
            self.cancelled.store(true, Ordering::Release);
        }
        self.wait_for_outstanding(started_at, generation_end)?;
        let finished_at = if has_drain_work {
            started_at.elapsed()
        } else {
            generation_end_elapsed
        };
        self.metrics.finish(finished_at)
    }

    fn process_ready_completions(
        &mut self,
        started_at: Instant,
        generation_end: Duration,
    ) -> Result<usize> {
        let mut processed = 0;
        while processed < MAX_READY_COMPLETIONS_PER_TICK {
            let now = started_at.elapsed();
            if now >= generation_end || now >= self.metrics.next_capture_deadline() {
                return Ok(processed);
            }
            match self.completion_rx.try_recv() {
                Ok(completion) => {
                    let now = started_at.elapsed();
                    if now >= generation_end || now >= self.metrics.next_capture_deadline() {
                        self.defer_completion(completion);
                        return Ok(processed);
                    }
                    self.process_completion(completion, now, generation_end)?;
                    processed += 1;
                }
                Err(TryRecvError::Empty) => return Ok(processed),
                Err(TryRecvError::Disconnected) => {
                    anyhow::bail!("workload completion channel disconnected unexpectedly")
                }
            }
        }
        Ok(processed)
    }

    fn process_completion(
        &mut self,
        completion: TaskCompletion,
        now: Duration,
        generation_end: Duration,
    ) -> Result<()> {
        let should_replace = {
            let runtime = self
                .runtimes
                .get_mut(completion.runtime_id)
                .context("completion references an unknown task runtime")?;
            ensure!(runtime.outstanding > 0, "task runtime completion underflow");
            runtime.outstanding -= 1;
            runtime.config.arrival_mode == ArrivalMode::ClosedLoop
                && now < generation_end
                && now < runtime.end_time
        };
        self.metrics.record_completion(
            completion.runtime_id,
            completion.completed,
            completion.latency,
        );
        if let Some(slow_task) = completion.slow_task {
            self.metrics.record_slow_task(slow_task);
        }

        if should_replace {
            self.metrics.record_offered(completion.runtime_id, 1);
            self.submit_next(completion.runtime_id)?;
        }
        Ok(())
    }

    fn process_due_events(&mut self, started_at: Instant, generation_end: Duration) -> Result<()> {
        while let Some(Reverse(event)) = self.events.peek().copied() {
            let now = started_at.elapsed();
            if now >= generation_end || now >= self.metrics.next_capture_deadline() {
                break;
            }
            if event.at > now {
                break;
            }
            self.events.pop();
            if now >= self.runtimes[event.runtime_id].end_time {
                self.suppress_remaining_open_loop_arrivals(event.runtime_id);
                continue;
            }
            match self.runtimes[event.runtime_id].config.arrival_mode {
                ArrivalMode::ClosedLoop => {
                    if self.fill_closed_loop(event.runtime_id, started_at, generation_end)? {
                        self.events.push(Reverse(GeneratorEvent {
                            at: started_at.elapsed(),
                            runtime_id: event.runtime_id,
                        }));
                        break;
                    }
                }
                ArrivalMode::OpenLoop => {
                    self.process_open_loop(event.runtime_id, started_at, generation_end, now)?
                }
            }
        }
        Ok(())
    }

    fn fill_closed_loop(
        &mut self,
        runtime_id: usize,
        started_at: Instant,
        generation_end: Duration,
    ) -> Result<bool> {
        let available = {
            let runtime = &self.runtimes[runtime_id];
            runtime.config.concurrency - runtime.outstanding
        };
        for _ in 0..available {
            let now = started_at.elapsed();
            if now >= generation_end || now >= self.runtimes[runtime_id].end_time {
                break;
            }
            if now >= self.metrics.next_capture_deadline() {
                return Ok(true);
            }
            self.metrics.record_offered(runtime_id, 1);
            self.submit_next(runtime_id)?;
        }
        Ok(false)
    }

    fn process_open_loop(
        &mut self,
        runtime_id: usize,
        started_at: Instant,
        generation_end: Duration,
        now: Duration,
    ) -> Result<()> {
        let (start_sequence, due, available) = {
            let runtime = &self.runtimes[runtime_id];
            (
                runtime.next_offer_sequence,
                runtime.arrivals_due(now).max(runtime.next_offer_sequence),
                runtime.config.concurrency - runtime.outstanding,
            )
        };
        let pending = due.saturating_sub(start_sequence);
        let to_submit = pending.min(available as u64);
        let mut submitted = 0;
        let mut interrupted = false;
        for _ in 0..to_submit {
            let current = started_at.elapsed();
            if current >= generation_end
                || current >= self.runtimes[runtime_id].end_time
                || current >= self.metrics.next_capture_deadline()
            {
                interrupted = true;
                break;
            }
            self.submit_next(runtime_id)?;
            submitted += 1;
        }
        let processed = if interrupted { submitted } else { pending };
        self.metrics.record_offered(runtime_id, processed);
        self.metrics
            .record_suppressed(runtime_id, processed.saturating_sub(submitted));

        // Arrivals beyond the concurrency limit are intentionally suppressed.
        // If a metrics boundary interrupts submission, leave the remaining due
        // arrivals on the timeline so the controller resumes them after the
        // snapshot.
        self.runtimes[runtime_id].next_offer_sequence = start_sequence + processed;
        let runtime = &self.runtimes[runtime_id];
        if runtime.next_offer_sequence < runtime.total_arrivals {
            let next_at = runtime.arrival_deadline(runtime.next_offer_sequence);
            self.events.push(Reverse(GeneratorEvent {
                at: next_at,
                runtime_id,
            }));
        }
        Ok(())
    }

    fn submit_next(&mut self, runtime_id: usize) -> Result<()> {
        let (identity, keyspace_id, priority, task_config) = {
            let runtime = &mut self.runtimes[runtime_id];
            let identity = TaskIdentity {
                seed: self.config.random_seed,
                runtime_id: runtime.sample_identity_id,
                offer_sequence: runtime.next_offer_sequence,
            };
            runtime.next_offer_sequence = runtime
                .next_offer_sequence
                .checked_add(1)
                .context("task offer sequence overflow")?;
            (
                identity,
                runtime.keyspace_id,
                command_priority(runtime.config.priority),
                runtime.config.clone(),
            )
        };

        let simulation_started_at = self
            .simulation_started_at
            .context("workload simulation start time is not initialized")?;
        let completion_tx = self.completion_tx.clone();
        let cancelled = self.cancelled.clone();
        let submitted_at = Instant::now();
        let admission = &self.config.read_pool.fair_global_cpu_admission;
        let fair_global_cpu_request_deadline = admission
            .enabled
            .then(|| {
                submitted_at
                    .checked_add(admission.request_timeout.0)
                    .context("fair global CPU admission request deadline overflow")
            })
            .transpose()?;
        let admit_interval_us = admission.enabled.then_some(admission.admit_interval_us);
        let task_id = self.next_task_id;
        self.next_task_id = self
            .next_task_id
            .checked_add(1)
            .context("YATP task ID overflow")?;
        let task_poll_metrics = self.metrics.task_poll_metrics(runtime_id);
        let TaskObservers {
            poll_tracker,
            result_handle,
        } = TaskObservers::new(
            task_id,
            runtime_id,
            identity.offer_sequence,
            simulation_started_at,
            submitted_at,
            task_config.slow_threshold.0,
            completion_tx,
            task_poll_metrics,
        );
        let poll_tracker = poll_tracker as Box<dyn YatpPollEventTracker + Send>;
        let mut task_options =
            ReadPoolTaskOptions::default().with_poll_tracker_observer(poll_tracker);
        if let Some(deadline) = fair_global_cpu_request_deadline {
            task_options =
                task_options.with_fair_cpu_admission(FairCpuTaskType::DagSelect, Some(deadline));
        }
        let spawn_result = self.handle.spawn_with_options(
            async move {
                let completed = run_task(task_config, identity, cancelled, admit_interval_us).await;
                result_handle.record(completed);
            },
            priority,
            keyspace_id,
            task_id,
            task_options,
        );
        if let Err(error) = spawn_result {
            return Err(error).context("failed to submit task to YATP");
        }
        self.runtimes[runtime_id].outstanding += 1;
        self.metrics.record_submitted(runtime_id);
        Ok(())
    }

    fn suppress_unprocessed_open_loop_arrivals(&mut self) {
        for runtime_id in 0..self.runtimes.len() {
            self.suppress_remaining_open_loop_arrivals(runtime_id);
        }
    }

    fn suppress_remaining_open_loop_arrivals(&mut self, runtime_id: usize) {
        let runtime = &mut self.runtimes[runtime_id];
        if runtime.config.arrival_mode != ArrivalMode::OpenLoop {
            return;
        }
        let unprocessed = runtime
            .total_arrivals
            .saturating_sub(runtime.next_offer_sequence);
        if unprocessed == 0 {
            return;
        }
        self.metrics.record_offered(runtime.runtime_id, unprocessed);
        self.metrics
            .record_suppressed(runtime.runtime_id, unprocessed);
        runtime.next_offer_sequence = runtime.total_arrivals;
    }

    fn wait_for_outstanding(
        &mut self,
        started_at: Instant,
        generation_end: Duration,
    ) -> Result<()> {
        let mut remaining = self.outstanding_task_count();
        while remaining > 0 {
            let now = started_at.elapsed();
            self.metrics.capture_if_due(now)?;
            if let Some(completion) = self.deferred_completion.take() {
                self.process_completion(completion, started_at.elapsed(), generation_end)?;
                remaining -= 1;
                continue;
            }
            let wait = self
                .metrics
                .next_capture_deadline()
                .saturating_sub(started_at.elapsed());
            match self.completion_rx.recv_timeout(wait) {
                Ok(completion) => {
                    let now = started_at.elapsed();
                    if now >= self.metrics.next_capture_deadline() {
                        self.defer_completion(completion);
                    } else {
                        self.process_completion(completion, now, generation_end)?;
                        remaining -= 1;
                    }
                }
                Err(RecvTimeoutError::Timeout) => {}
                Err(RecvTimeoutError::Disconnected) => {
                    anyhow::bail!("workload completion channel disconnected while draining")
                }
            }
        }
        Ok(())
    }

    fn defer_completion(&mut self, completion: TaskCompletion) {
        debug_assert!(self.deferred_completion.is_none());
        self.deferred_completion = Some(completion);
    }

    fn outstanding_task_count(&self) -> usize {
        self.runtimes
            .iter()
            .map(|runtime| runtime.outstanding)
            .sum()
    }
}

fn build_workload_plan(config: &SimulatorConfig) -> Result<WorkloadPlan> {
    let mut keyspaces = Vec::new();
    let mut tasks = Vec::new();
    let mut runtimes = Vec::new();
    let mut next_keyspace_id = 1u64;
    let mut next_sample_identity_id = 0u64;
    for group in &config.keyspace_groups {
        for group_instance in 1..=group.count {
            let keyspace_id = u32::try_from(next_keyspace_id).context("keyspace ID overflow")?;
            let keyspace_name = if group.count == 1 {
                group.name.clone()
            } else {
                format!("{}-{group_instance}", group.name)
            };
            keyspaces.push(KeyspaceDescriptor {
                keyspace_id,
                keyspace_name: keyspace_name.clone(),
                group_name: group.name.clone(),
                group_instance,
            });
            for task in &group.tasks {
                let sample_identity_id = next_sample_identity_id;
                next_sample_identity_id = next_sample_identity_id
                    .checked_add(1)
                    .context("task sample identity overflow")?;
                if !task.enabled {
                    continue;
                }
                let runtime_id = runtimes.len();
                let end_time = config.task_end_time(task)?;
                let total_arrivals = match task.arrival_mode {
                    ArrivalMode::OpenLoop => {
                        let active_seconds =
                            end_time.saturating_sub(task.start_time.0).as_secs_f64();
                        let qps = task
                            .qps
                            .context("validated open-loop task should have an arrival rate")?;
                        (active_seconds * qps).ceil().max(1.0) as u64
                    }
                    ArrivalMode::ClosedLoop => 0,
                };
                tasks.push(TaskDescriptor {
                    runtime_id,
                    keyspace_id,
                    keyspace_name: keyspace_name.clone(),
                    group_name: group.name.clone(),
                    group_instance,
                    task_name: task.name.clone(),
                });
                runtimes.push(TaskRuntime {
                    runtime_id,
                    sample_identity_id,
                    keyspace_id,
                    config: Arc::new(task.clone()),
                    start_time: task.start_time.0,
                    end_time,
                    total_arrivals,
                    outstanding: 0,
                    next_offer_sequence: 0,
                });
            }
            next_keyspace_id += 1;
        }
    }
    Ok(WorkloadPlan {
        keyspaces,
        tasks,
        runtimes,
    })
}

#[derive(Clone, Copy, Debug)]
struct TaskIdentity {
    seed: u64,
    runtime_id: u64,
    offer_sequence: u64,
}

#[derive(Clone, Copy)]
#[repr(u64)]
enum SampleKind {
    TaskCpu = 0x7a52_b14d_9f6c_03e1,
    SliceCpu = 0x3df1_05a9_c627_84be,
    BlockingWall = 0xc4e8_6913_2dab_705f,
}

async fn run_task(
    config: Arc<TaskConfig>,
    identity: TaskIdentity,
    cancelled: Arc<AtomicBool>,
    admit_interval_us: Option<u64>,
) -> bool {
    let admit_interval = admit_interval_us.map(Duration::from_micros);
    let mut sampled_task_cpu_target = None;
    let mut consumed_task_cpu = Duration::ZERO;
    let mut current_slice = None;
    let mut logical_slice_index = 0u64;

    loop {
        let (action, poll_cpu, task_cpu_target) = {
            // This timer is intentionally created before task and slice setup
            // so their CPU is part of the configured task and slice targets.
            // The scope ends before any await, so it always measures one worker
            // thread and the future remains Send.
            let poll_started_at = ThreadTime::now();
            let task_cpu_target = *sampled_task_cpu_target.get_or_insert_with(|| {
                Duration::from_micros(sample_duration_us(
                    config.task_cpu_us,
                    config.task_cpu_jitter_us,
                    1,
                    identity,
                    0,
                    SampleKind::TaskCpu,
                ))
            });
            let remaining_task_cpu = task_cpu_target.saturating_sub(consumed_task_cpu);
            debug_assert!(remaining_task_cpu > Duration::ZERO);
            let slice = current_slice.get_or_insert_with(|| {
                LogicalSlice::new(
                    Duration::from_micros(sample_duration_us(
                        config.slice_cpu_us,
                        config.slice_cpu_jitter_us,
                        1,
                        identity,
                        logical_slice_index,
                        SampleKind::SliceCpu,
                    )),
                    Duration::from_micros(sample_duration_us(
                        config.blocking_wall_us,
                        config.blocking_wall_jitter_us,
                        0,
                        identity,
                        logical_slice_index,
                        SampleKind::BlockingWall,
                    )),
                )
            });
            slice.cpu_target = slice.cpu_target.min(remaining_task_cpu);

            let action = run_slice(&poll_started_at, slice, &cancelled, admit_interval, || {
                let mut now = tikv_util::time::LazyInstant::default();
                fair_cpu_checkpoint(&mut now)
            });
            (action, poll_started_at.elapsed(), task_cpu_target)
        };
        if action == SliceRunAction::Cancelled {
            return false;
        }
        consumed_task_cpu = consumed_task_cpu.saturating_add(poll_cpu);
        if consumed_task_cpu >= task_cpu_target {
            return true;
        }
        if cancelled.load(Ordering::Acquire) {
            return false;
        }

        match action {
            SliceRunAction::YieldToScheduler => {
                yatp::task::future::yield_to_scheduler().await;
            }
            SliceRunAction::Reschedule => {
                current_slice = None;
                logical_slice_index = logical_slice_index.wrapping_add(1);
                yatp::task::future::reschedule().await;
            }
            SliceRunAction::Cancelled => {
                unreachable!("cancelled slices return before task CPU completion is evaluated")
            }
        }
    }
}

#[derive(Debug)]
struct LogicalSlice {
    cpu_target: Duration,
    blocking_wall: Option<Duration>,
}

impl LogicalSlice {
    fn new(cpu_target: Duration, blocking_wall: Duration) -> Self {
        debug_assert!(cpu_target > Duration::ZERO);
        Self {
            cpu_target,
            blocking_wall: (blocking_wall > Duration::ZERO).then_some(blocking_wall),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum SliceRunAction {
    YieldToScheduler,
    Reschedule,
    Cancelled,
}

/// Runs the synchronous part of one task poll without yielding.
///
/// CPU is measured from `poll_started_at`, which the caller creates before any
/// per-poll task setup. Returning [`SliceRunAction::YieldToScheduler`] retains
/// the logical slice so the next poll restarts its complete sampled CPU target.
/// Returning [`SliceRunAction::Reschedule`] finishes the logical slice.
/// Returning [`SliceRunAction::Cancelled`] means cancellation was observed
/// before the slice completed its simulated CPU work.
fn run_slice(
    poll_started_at: &ThreadTime,
    slice: &mut LogicalSlice,
    cancelled: &AtomicBool,
    admit_interval: Option<Duration>,
    mut checkpoint: impl FnMut() -> FairCpuBudgetAction,
) -> SliceRunAction {
    debug_assert!(slice.cpu_target > Duration::ZERO);

    if cancelled.load(Ordering::Acquire) {
        return SliceRunAction::Cancelled;
    }

    if let Some(blocking_wall) = slice.blocking_wall.take() {
        thread::sleep(blocking_wall);
    }

    if cancelled.load(Ordering::Acquire) {
        return SliceRunAction::Cancelled;
    }

    let Some(interval) = admit_interval else {
        consume_cpu_until(poll_started_at, slice.cpu_target);
        return SliceRunAction::Reschedule;
    };

    let mut schedule = CpuCheckpointSchedule::new(slice.cpu_target, interval);
    loop {
        consume_cpu_until(poll_started_at, schedule.deadline());
        if checkpoint() == FairCpuBudgetAction::YieldToScheduler {
            return SliceRunAction::YieldToScheduler;
        }
        if cancelled.load(Ordering::Acquire) {
            return SliceRunAction::Cancelled;
        }
        if poll_started_at.elapsed() >= slice.cpu_target {
            return SliceRunAction::Reschedule;
        }
        schedule.advance(poll_started_at.elapsed());
    }
}

#[derive(Debug)]
struct CpuCheckpointSchedule {
    cpu_limit: Duration,
    interval: Duration,
    next_deadline: Duration,
}

impl CpuCheckpointSchedule {
    fn new(cpu_limit: Duration, interval: Duration) -> Self {
        debug_assert!(cpu_limit > Duration::ZERO);
        debug_assert!(interval > Duration::ZERO);
        Self {
            cpu_limit,
            interval,
            next_deadline: interval.min(cpu_limit),
        }
    }

    fn deadline(&self) -> Duration {
        self.next_deadline
    }

    fn advance(&mut self, elapsed: Duration) {
        debug_assert!(self.next_deadline < self.cpu_limit);
        let nominal_deadline = self.next_deadline.saturating_add(self.interval);
        self.next_deadline = if nominal_deadline <= elapsed {
            elapsed.saturating_add(self.interval)
        } else {
            nominal_deadline
        }
        .min(self.cpu_limit);
    }
}

fn consume_cpu_until(started_at: &ThreadTime, target_elapsed: Duration) {
    while started_at.elapsed() < target_elapsed {
        for _ in 0..CPU_TIME_CHECK_BATCH_SIZE {
            std::hint::black_box(());
        }
    }
}

fn sample_duration_us(
    base: u64,
    jitter: u64,
    minimum: u64,
    identity: TaskIdentity,
    slice_index: u64,
    kind: SampleKind,
) -> u64 {
    if jitter == 0 {
        return base.max(minimum);
    }
    let lower = base.saturating_sub(jitter).max(minimum);
    let upper = base
        .checked_add(jitter)
        .expect("jitter upper bound should be validated");
    let mut rng = StdRng::seed_from_u64(sample_seed(identity, slice_index, kind));
    rng.gen_range(lower..=upper)
}

fn sample_seed(identity: TaskIdentity, slice_index: u64, kind: SampleKind) -> u64 {
    let mut state = identity.seed;
    state = splitmix64(state ^ identity.runtime_id);
    state = splitmix64(state ^ identity.offer_sequence);
    state = splitmix64(state ^ slice_index);
    splitmix64(state ^ kind as u64)
}

fn splitmix64(mut value: u64) -> u64 {
    value = value.wrapping_add(0x9e37_79b9_7f4a_7c15);
    value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
    value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
    value ^ (value >> 31)
}

fn command_priority(priority: Priority) -> CommandPri {
    match priority {
        Priority::High => CommandPri::High,
        Priority::Normal => CommandPri::Normal,
        Priority::Low => CommandPri::Low,
    }
}

#[cfg(test)]
mod tests {
    use std::{
        future::Future,
        task::{Context, Poll, Waker},
    };

    use super::*;
    use crate::{config::SimulatorConfig, metrics::KeyspaceMetricStore, pool::SimulatorReadPool};

    #[test]
    fn test_cpu_checkpoint_schedule_includes_final_remainder() {
        let mut schedule =
            CpuCheckpointSchedule::new(Duration::from_micros(120), Duration::from_micros(50));
        assert_eq!(schedule.deadline(), Duration::from_micros(50));
        schedule.advance(Duration::from_micros(50));
        assert_eq!(schedule.deadline(), Duration::from_micros(100));
        schedule.advance(Duration::from_micros(100));
        assert_eq!(schedule.deadline(), Duration::from_micros(120));
    }

    #[test]
    fn test_cpu_checkpoint_schedule_catches_up_after_overshoot() {
        let mut schedule =
            CpuCheckpointSchedule::new(Duration::from_micros(200), Duration::from_micros(50));
        schedule.advance(Duration::from_micros(80));
        assert_eq!(schedule.deadline(), Duration::from_micros(100));
        schedule.advance(Duration::from_micros(130));
        assert_eq!(schedule.deadline(), Duration::from_micros(150));
        schedule.advance(Duration::from_micros(180));
        assert_eq!(schedule.deadline(), Duration::from_micros(200));
    }

    #[test]
    fn test_run_slice_returns_reschedule_after_actual_cpu_target() {
        let cancelled = AtomicBool::new(false);
        let mut slice = LogicalSlice::new(Duration::from_millis(2), Duration::ZERO);
        let poll_started_at = ThreadTime::now();
        let action = run_slice(&poll_started_at, &mut slice, &cancelled, None, || {
            unreachable!("admission is disabled")
        });
        let poll_cpu = poll_started_at.elapsed();
        assert_eq!(action, SliceRunAction::Reschedule);
        assert!(poll_cpu >= Duration::from_millis(2));
    }

    #[test]
    fn test_run_slice_admission_yield_consumes_wall_once() {
        let cancelled = AtomicBool::new(false);
        let mut slice = LogicalSlice::new(Duration::from_millis(100), Duration::from_millis(1));
        let poll_started_at = ThreadTime::now();
        let action = run_slice(
            &poll_started_at,
            &mut slice,
            &cancelled,
            Some(Duration::from_millis(1)),
            || FairCpuBudgetAction::YieldToScheduler,
        );
        assert_eq!(action, SliceRunAction::YieldToScheduler);
        assert!(poll_started_at.elapsed() >= Duration::from_millis(1));
        assert_eq!(slice.cpu_target, Duration::from_millis(100));
        assert_eq!(slice.blocking_wall, None);

        cancelled.store(true, Ordering::Release);
        let resumed_at = ThreadTime::now();
        let resumed = run_slice(
            &resumed_at,
            &mut slice,
            &cancelled,
            Some(Duration::from_millis(1)),
            || unreachable!("cancelled slices do not checkpoint"),
        );
        assert_eq!(resumed, SliceRunAction::Cancelled);
    }

    #[test]
    fn test_run_task_reports_cancelled_before_poll_cpu_can_complete_target() {
        let mut config: SimulatorConfig = toml::from_str(
            r#"
duration = "1s"
stats-print-interval = "1s"
cooldown-duration = "0s"
report-path = "unused.html"
summary-path = "unused.json"

[read-pool]
workers = 1

[[keyspace-groups]]
name = "test"

[[keyspace-groups.tasks]]
name = "cancelled"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 1
slice-cpu-us = 1
"#,
        )
        .unwrap();
        let task = Arc::new(config.keyspace_groups[0].tasks.remove(0));
        let cancelled = Arc::new(AtomicBool::new(true));
        let mut future = std::pin::pin!(run_task(
            task,
            TaskIdentity {
                seed: config.random_seed,
                runtime_id: 0,
                offer_sequence: 0,
            },
            cancelled,
            None,
        ));
        let mut context = Context::from_waker(Waker::noop());

        assert_eq!(future.as_mut().poll(&mut context), Poll::Ready(false));
    }

    #[test]
    fn test_workload_collects_slow_tasks_and_keyspace_queue_metrics() {
        let config: SimulatorConfig = toml::from_str(
            r#"
duration = "20ms"
stats-print-interval = "100ms"
cooldown-duration = "0s"
report-path = "unused.html"
summary-path = "unused.json"

[read-pool]
workers = 1

[read-pool.keyspace-fairness]
enabled = true

[[keyspace-groups]]
name = "test"

[[keyspace-groups.tasks]]
name = "slow"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 200
slice-cpu-us = 50
slow-threshold = "1us"

[[keyspace-groups.tasks]]
name = "fast"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 100
slice-cpu-us = 50
"#,
        )
        .unwrap();
        let metric_store = KeyspaceMetricStore::new().unwrap();
        let pool =
            SimulatorReadPool::new(&config.read_pool, metric_store.keyspace_metrics()).unwrap();
        let handle = pool.handle();
        let worker_cpu_time = pool.worker_cpu_time_source();

        let report = run(&config, &handle, metric_store, worker_cpu_time).unwrap();
        drop(handle);
        pool.shutdown();

        assert!(!report.slow_tasks.is_empty());
        assert!(report.slow_tasks.iter().all(|task| {
            task.runtime_id == 0
                && task.latency_seconds > task.slow_threshold_seconds
                && !task.polls.is_empty()
        }));
        assert_eq!(report.task_runtimes.len(), 2);
        assert_eq!(report.summary.total.keyspaces.len(), 1);
        assert!(
            report.summary.total.keyspaces[0]
                .fair_queue_length
                .as_ref()
                .is_some_and(|queue_length| queue_length.count > 0)
        );
        assert!(
            report
                .slow_tasks
                .windows(2)
                .all(|tasks| { tasks[0].submitted_at_seconds <= tasks[1].submitted_at_seconds })
        );
    }

    #[test]
    fn test_jitter_is_deterministic_and_bounded() {
        let identity = TaskIdentity {
            seed: 7,
            runtime_id: 2,
            offer_sequence: 5,
        };
        let first = sample_duration_us(1_000, 200, 1, identity, 3, SampleKind::SliceCpu);
        let second = sample_duration_us(1_000, 200, 1, identity, 3, SampleKind::SliceCpu);
        assert_eq!(first, second);
        assert!((800..=1_200).contains(&first));
        assert_eq!(
            sample_duration_us(1_000, 0, 1, identity, 3, SampleKind::SliceCpu),
            1_000
        );
    }

    #[test]
    fn test_blocking_jitter_can_reach_zero() {
        let identity = TaskIdentity {
            seed: 11,
            runtime_id: 0,
            offer_sequence: 0,
        };
        let sampled = sample_duration_us(10, 20, 0, identity, 0, SampleKind::BlockingWall);
        assert!(sampled <= 30);
    }

    #[test]
    fn test_open_loop_deadlines_are_half_open() {
        let task = Arc::new(TaskConfig {
            name: "read".to_owned(),
            enabled: true,
            arrival_mode: ArrivalMode::OpenLoop,
            qps: Some(2.0),
            concurrency: 1,
            task_cpu_us: 1,
            task_cpu_jitter_us: 0,
            slice_cpu_us: 1,
            slice_cpu_jitter_us: 0,
            blocking_wall_us: 0,
            blocking_wall_jitter_us: 0,
            priority: Priority::Normal,
            start_time: tikv_util::config::ReadableDuration::ZERO,
            duration: tikv_util::config::ReadableDuration::secs(1),
            slow_threshold: tikv_util::config::ReadableDuration::ZERO,
        });
        let runtime = TaskRuntime {
            runtime_id: 0,
            sample_identity_id: 0,
            keyspace_id: 1,
            config: task,
            start_time: Duration::ZERO,
            end_time: Duration::from_secs(1),
            total_arrivals: 2,
            outstanding: 0,
            next_offer_sequence: 0,
        };
        assert_eq!(runtime.arrival_deadline(0), Duration::ZERO);
        assert_eq!(runtime.arrival_deadline(1), Duration::from_millis(500));
        assert_eq!(runtime.total_arrivals, 2);
        assert!(runtime.arrival_deadline(runtime.total_arrivals - 1) < runtime.end_time);
    }

    #[test]
    fn test_runtime_expansion_uses_one_id_per_keyspace() {
        let config: SimulatorConfig = toml::from_str(
            r#"
duration = "1s"

[[keyspace-groups]]
name = "group"
count = 2

[[keyspace-groups.tasks]]
name = "read"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5

[[keyspace-groups.tasks]]
name = "scan"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5

[[keyspace-groups.tasks]]
name = "disabled"
enabled = false
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5
"#,
        )
        .unwrap();
        let plan = build_workload_plan(&config).unwrap();
        assert_eq!(plan.keyspaces.len(), 2);
        assert_eq!(plan.tasks.len(), 4);
        assert_eq!(plan.runtimes.len(), 4);
        assert_eq!(plan.runtimes[0].keyspace_id, 1);
        assert_eq!(plan.runtimes[1].keyspace_id, 1);
        assert_eq!(plan.runtimes[2].keyspace_id, 2);
        assert_eq!(plan.runtimes[3].keyspace_id, 2);
        assert_eq!(
            plan.runtimes
                .iter()
                .map(|runtime| runtime.sample_identity_id)
                .collect::<Vec<_>>(),
            vec![0, 1, 3, 4]
        );
        assert_eq!(plan.keyspaces[0].keyspace_name, "group-1");
        assert_eq!(plan.keyspaces[0].group_instance, 1);
        assert_eq!(plan.keyspaces[1].keyspace_name, "group-2");
        for (runtime_id, task) in plan.tasks.iter().enumerate() {
            assert_eq!(task.runtime_id, runtime_id);
            assert_eq!(task.keyspace_id, plan.runtimes[runtime_id].keyspace_id);
            assert_ne!(task.task_name, "disabled");
        }
    }

    #[test]
    fn test_workload_plan_keeps_idle_keyspaces() {
        let config: SimulatorConfig = toml::from_str(
            r#"
duration = "1s"

[[keyspace-groups]]
name = "active"

[[keyspace-groups.tasks]]
name = "read"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5

[[keyspace-groups]]
name = "idle"
count = 2
"#,
        )
        .unwrap();
        let plan = build_workload_plan(&config).unwrap();
        assert_eq!(plan.keyspaces.len(), 3);
        assert_eq!(plan.tasks.len(), 1);
        assert_eq!(plan.runtimes.len(), 1);
        assert_eq!(plan.keyspaces[0].keyspace_name, "active");
        assert_eq!(plan.keyspaces[1].keyspace_name, "idle-1");
        assert_eq!(plan.keyspaces[2].keyspace_name, "idle-2");
    }
}