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

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt::Write as _,
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use anyhow::{Context, Result, ensure};
use prometheus::{
    Counter, CounterVec, Histogram, HistogramOpts, HistogramVec, IntGauge, core::Collector,
    exponential_buckets,
};
use serde::Serialize;
use tikv::read_pool::YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS;
use tikv_util::yatp_pool::{KeyspaceInnerMetrics, KeyspaceMetrics};

use crate::config::SimulatorConfig;

pub(crate) trait WorkerCpuTimeSource: Send + Sync {
    /// Returns cumulative CPU time consumed by all YATP worker threads.
    fn snapshot_cpu_nanos(&self) -> Result<u64>;
}

const QUEUE_BUCKET_START: f64 = 1.0;
const QUEUE_BUCKET_FACTOR: f64 = 2.0;
// Keep queue-length and waited-slice histograms aligned with the production
// read-pool metrics in `src/read_pool.rs`.
const QUEUE_BUCKET_COUNT: usize = 12;
const WAIT_BUCKET_START_SECONDS: f64 = 1e-5;
const WAIT_BUCKET_FACTOR: f64 = 4.0;
const WAIT_BUCKET_COUNT: usize = 12;
// Whole-task latency needs finer buckets around the expected millisecond range
// than internal queue-wait metrics. In particular, keep 5 ms steps from 10 ms
// through 50 ms so a small tail does not jump directly from 10.24 ms to
// 40.96 ms, while retaining coarse coverage for unusually long tasks.
const TASK_LATENCY_BUCKETS_SECONDS: &[f64] = &[
    0.000_01, 0.000_04, 0.000_16, 0.000_64, 0.002_56, 0.005, 0.010, 0.015, 0.020, 0.025, 0.030,
    0.035, 0.040, 0.045, 0.050, 0.075, 0.100, 0.160, 0.250, 0.500, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0,
    120.0, 300.0, 600.0,
];

/// Owns the unregistered production metrics installed on the simulated read
/// pool. Reports read only the production fair-queue length histogram; all
/// other report metrics remain simulator-owned at task-runtime scope.
pub(crate) struct KeyspaceMetricStore {
    queue_length: HistogramVec,
    keyspace_metrics: Arc<KeyspaceMetrics>,
}

impl KeyspaceMetricStore {
    pub(crate) fn new() -> Result<Self> {
        let queue_length = HistogramVec::new(
            HistogramOpts::new(
                "yatp_sim_keyspace_queue_length",
                "Per-keyspace queued task slices before a task slice is pushed",
            )
            .buckets(exponential_buckets(
                QUEUE_BUCKET_START,
                QUEUE_BUCKET_FACTOR,
                QUEUE_BUCKET_COUNT,
            )?),
            &["keyspace_id"],
        )?;
        let queue_task_slice_waits = HistogramVec::new(
            HistogramOpts::new(
                "yatp_sim_keyspace_task_slice_waits",
                "Successful fair-queue pops while a task slice waits",
            )
            .buckets(exponential_buckets(
                QUEUE_BUCKET_START,
                QUEUE_BUCKET_FACTOR,
                QUEUE_BUCKET_COUNT,
            )?),
            &["keyspace_id"],
        )?;
        let queue_task_slice_wait_duration = HistogramVec::new(
            HistogramOpts::new(
                "yatp_sim_keyspace_task_slice_wait_duration_seconds",
                "Task-slice wait duration in the keyspace fair queue",
            )
            .buckets(YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()),
            &["keyspace_id"],
        )?;
        let task_slice_schedule_interval = HistogramVec::new(
            HistogramOpts::new(
                "yatp_sim_keyspace_task_slice_schedule_interval_seconds",
                "Interval between consecutive scheduled task slices",
            )
            .buckets(YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()),
            &["keyspace_id"],
        )?;
        let task_slice_cpu_duration = HistogramVec::new(
            HistogramOpts::new(
                "yatp_sim_keyspace_task_slice_cpu_duration_seconds",
                "CPU duration of each YATP task slice for each keyspace",
            )
            .buckets(YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()),
            &["keyspace_id"],
        )?;
        let task_wall_time = CounterVec::new(
            prometheus::Opts::new(
                "yatp_sim_keyspace_task_wall_seconds_total",
                "Wall time consumed by YATP task polls for each keyspace",
            ),
            &["keyspace_id"],
        )?;
        let fair_queue_keyspace_count = IntGauge::new(
            "yatp_sim_keyspace_fair_queue_group_count",
            "Current number of groups in the keyspace fair queue",
        )?;

        let create_queue_length = queue_length.clone();
        let keyspace_metrics = KeyspaceMetrics::new(
            fair_queue_keyspace_count,
            move |keyspace_id| {
                let keyspace_id = keyspace_id.to_string();
                let labels = [keyspace_id.as_str()];
                KeyspaceInnerMetrics::new(
                    create_queue_length.with_label_values(&labels),
                    queue_task_slice_waits.with_label_values(&labels),
                    queue_task_slice_wait_duration.with_label_values(&labels),
                    task_slice_schedule_interval.with_label_values(&labels),
                    task_slice_cpu_duration.with_label_values(&labels),
                    task_wall_time.with_label_values(&labels),
                )
            },
            |_| {},
        );

        Ok(Self {
            queue_length,
            keyspace_metrics,
        })
    }

    pub(crate) fn keyspace_metrics(&self) -> Arc<KeyspaceMetrics> {
        self.keyspace_metrics.clone()
    }

    fn snapshot_queue_lengths(
        &self,
        keyspace_ids: &[u32],
        fairness_enabled: bool,
    ) -> BTreeMap<u32, HistogramSnapshot> {
        if !fairness_enabled {
            return BTreeMap::new();
        }
        keyspace_ids
            .iter()
            .map(|keyspace_id| {
                let keyspace_id_label = keyspace_id.to_string();
                let labels = [keyspace_id_label.as_str()];
                (
                    *keyspace_id,
                    HistogramSnapshot::from_prometheus(
                        &self.queue_length.with_label_values(&labels),
                    ),
                )
            })
            .collect()
    }
}

#[derive(Clone, Debug, Serialize)]
pub struct KeyspaceDescriptor {
    pub keyspace_id: u32,
    pub keyspace_name: String,
    pub group_name: String,
    pub group_instance: usize,
}

#[derive(Clone, Debug, Serialize)]
pub struct TaskDescriptor {
    pub runtime_id: usize,
    pub keyspace_id: u32,
    pub keyspace_name: String,
    pub group_name: String,
    pub group_instance: usize,
    pub task_name: String,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum ReportPhase {
    Warmup,
    Measurement,
    Drain,
    Total,
    Workload,
}

#[derive(Clone, Debug, Serialize)]
pub struct HistogramBucketReport {
    pub upper_bound: f64,
    pub cumulative_count: u64,
}

#[derive(Clone, Debug, Serialize)]
pub struct DistributionReport {
    pub count: u64,
    pub sum: f64,
    pub average: Option<f64>,
    pub p50: Option<f64>,
    pub p80: Option<f64>,
    pub p95: Option<f64>,
    pub p99: Option<f64>,
    pub overflow_count: u64,
    pub buckets: Vec<HistogramBucketReport>,
}

#[derive(Clone, Debug, Serialize)]
pub struct TaskMetricsReport {
    pub runtime_id: usize,
    pub keyspace_id: u32,
    pub keyspace_name: String,
    pub group_name: String,
    pub group_instance: usize,
    pub task_name: String,
    pub offered: u64,
    pub submitted: u64,
    pub suppressed: u64,
    pub completed: u64,
    pub cancelled: u64,
    pub outstanding_at_end: usize,
    pub throughput_qps: f64,
    pub task_latency_seconds: DistributionReport,
    pub task_poll_cpu_seconds: f64,
    pub task_poll_wall_seconds: f64,
    pub task_poll_cpu_cores: f64,
    pub task_poll_wall_cores: f64,
    pub task_slice_wait_seconds: DistributionReport,
    pub task_slice_waits: DistributionReport,
}

#[derive(Clone, Debug, Serialize)]
pub struct KeyspaceMetricsReport {
    pub keyspace_id: u32,
    pub keyspace_name: String,
    pub group_name: String,
    pub group_instance: usize,
    pub offered: u64,
    pub submitted: u64,
    pub suppressed: u64,
    pub completed: u64,
    pub cancelled: u64,
    pub outstanding_at_end: usize,
    pub throughput_qps: f64,
    pub task_latency_seconds: DistributionReport,
    pub task_poll_cpu_seconds: f64,
    pub task_poll_wall_seconds: f64,
    pub task_poll_cpu_cores: f64,
    pub task_poll_wall_cores: f64,
    pub accounted_seconds: Option<f64>,
    pub accounted_service_share: Option<f64>,
    pub fair_queue_length: Option<DistributionReport>,
    pub fair_queue_task_slice_waits: Option<DistributionReport>,
    pub fair_queue_task_slice_wait_seconds: Option<DistributionReport>,
}

#[derive(Clone, Debug, Serialize)]
pub struct MetricsWindow {
    pub phase: ReportPhase,
    pub start_seconds: f64,
    pub end_seconds: f64,
    pub duration_seconds: f64,
    /// CPU consumed by all YATP worker threads during this window.
    pub total_cpu_seconds: f64,
    /// Average number of YATP worker CPU cores consumed during this window.
    pub total_cpu_cores: f64,
    pub keyspaces: Vec<KeyspaceMetricsReport>,
    pub tasks: Vec<TaskMetricsReport>,
}

#[derive(Clone, Debug, Serialize)]
pub struct RunSummary {
    pub warmup: Option<MetricsWindow>,
    pub measurement: MetricsWindow,
    pub drain: Option<MetricsWindow>,
    pub total: MetricsWindow,
    pub workload_phases: Vec<WorkloadPhaseReport>,
}

#[derive(Clone, Debug, Serialize)]
pub struct WorkloadPhaseReport {
    pub number: usize,
    pub active_runtime_ids: Vec<usize>,
    pub metrics: MetricsWindow,
}

/// Timing and scheduler details for one poll of a completed slow task.
#[derive(Clone, Debug, Serialize)]
pub struct SlowPollReport {
    /// One-based poll position within the task.
    pub poll_index: usize,
    /// One-based YATP dispatch position within the task.
    pub dispatch_index: usize,
    /// Whether this poll immediately followed another poll in the same
    /// dispatch.
    pub immediate_repoll: bool,
    /// Previous poll end, or task submission for the first poll, relative to
    /// simulation start.
    pub interval_start_seconds: f64,
    /// Actual poll start relative to simulation start.
    pub poll_start_seconds: f64,
    /// Actual poll end relative to simulation start.
    pub poll_end_seconds: f64,
    /// Total time between `interval_start_seconds` and `poll_start_seconds`.
    pub wait_seconds: f64,
    /// Portion of `wait_seconds` spent queued for a YATP worker.
    pub queue_wait_seconds: f64,
    /// Remaining pre-poll time not attributed to the YATP queue.
    pub other_wait_seconds: f64,
    /// Wall time spent inside this invocation of `Future::poll`.
    pub poll_wall_seconds: f64,
    /// Thread CPU time consumed inside this invocation of `Future::poll`.
    pub poll_cpu_seconds: f64,
    /// Fair-queue task slices dispatched while this dispatch was waiting.
    pub fair_queue_waited_task_slices: Option<u64>,
}

/// A completed task whose whole-task latency exceeded its configured threshold.
#[derive(Clone, Debug, Serialize)]
pub struct SlowTaskReport {
    /// Unique YATP task ID assigned by this simulator run.
    pub task_id: u64,
    /// Index into `ReportData::task_runtimes`.
    pub runtime_id: usize,
    /// Zero-based offered-task sequence within this runtime.
    pub offer_sequence: u64,
    /// Submission time relative to simulation start.
    pub submitted_at_seconds: f64,
    /// Final poll observation time relative to simulation start.
    pub completed_at_seconds: f64,
    /// Whole-task latency from submission through the final poll.
    pub latency_seconds: f64,
    /// Configured non-zero threshold exceeded by this task.
    pub slow_threshold_seconds: f64,
    /// Polls in execution order.
    pub polls: Vec<SlowPollReport>,
}

#[derive(Clone, Debug, Serialize)]
pub struct ReportData {
    pub schema_version: u32,
    pub generated_at_unix_seconds: u64,
    pub config: SimulatorConfig,
    pub keyspaces: Vec<KeyspaceDescriptor>,
    pub task_runtimes: Vec<TaskDescriptor>,
    pub windows: Vec<MetricsWindow>,
    pub summary: RunSummary,
    pub slow_tasks: Vec<SlowTaskReport>,
}

#[derive(Clone, Debug)]
struct HistogramSnapshot {
    count: u64,
    sum: f64,
    buckets: Vec<HistogramBucketSnapshot>,
}

#[derive(Clone, Debug)]
struct HistogramBucketSnapshot {
    upper_bound: f64,
    cumulative_count: u64,
}

impl HistogramSnapshot {
    fn from_prometheus(histogram: &Histogram) -> Self {
        let families = histogram.collect();
        let histogram = families[0].get_metric()[0].get_histogram();
        Self {
            count: histogram.get_sample_count(),
            sum: histogram.get_sample_sum(),
            buckets: histogram
                .get_bucket()
                .iter()
                .map(|bucket| HistogramBucketSnapshot {
                    upper_bound: bucket.get_upper_bound(),
                    cumulative_count: bucket.get_cumulative_count(),
                })
                .collect(),
        }
    }

    fn delta(&self, before: &Self) -> Self {
        debug_assert_eq!(self.buckets.len(), before.buckets.len());
        let buckets = self
            .buckets
            .iter()
            .zip(&before.buckets)
            .map(|(after, before)| {
                debug_assert_eq!(after.upper_bound, before.upper_bound);
                HistogramBucketSnapshot {
                    upper_bound: after.upper_bound,
                    cumulative_count: after
                        .cumulative_count
                        .saturating_sub(before.cumulative_count),
                }
            })
            .collect();
        Self {
            count: self.count.saturating_sub(before.count),
            sum: (self.sum - before.sum).max(0.0),
            buckets,
        }
    }

    fn merge(&mut self, other: &Self) {
        if self.buckets.is_empty() {
            *self = other.clone();
            return;
        }
        debug_assert_eq!(self.buckets.len(), other.buckets.len());
        self.count = self.count.saturating_add(other.count);
        self.sum += other.sum;
        for (target, source) in self.buckets.iter_mut().zip(&other.buckets) {
            debug_assert_eq!(target.upper_bound, source.upper_bound);
            target.cumulative_count = target
                .cumulative_count
                .saturating_add(source.cumulative_count);
        }
    }

    fn empty() -> Self {
        Self {
            count: 0,
            sum: 0.0,
            buckets: Vec::new(),
        }
    }

    fn report(&self) -> DistributionReport {
        let overflow_count = self
            .buckets
            .last()
            .map(|bucket| self.count.saturating_sub(bucket.cumulative_count))
            .unwrap_or(self.count);
        DistributionReport {
            count: self.count,
            sum: self.sum,
            average: (self.count > 0).then(|| self.sum / self.count as f64),
            p50: self.quantile_upper_bound(0.50),
            p80: self.quantile_upper_bound(0.80),
            p95: self.quantile_upper_bound(0.95),
            p99: self.quantile_upper_bound(0.99),
            overflow_count,
            buckets: self
                .buckets
                .iter()
                .map(|bucket| HistogramBucketReport {
                    upper_bound: bucket.upper_bound,
                    cumulative_count: bucket.cumulative_count,
                })
                .collect(),
        }
    }

    fn quantile_upper_bound(&self, quantile: f64) -> Option<f64> {
        if self.count == 0 {
            return None;
        }
        let rank = (self.count as f64 * quantile).ceil().max(1.0) as u64;
        self.buckets
            .iter()
            .find(|bucket| bucket.cumulative_count >= rank)
            .map(|bucket| bucket.upper_bound)
    }
}

#[derive(Clone, Debug)]
struct TaskMetricSnapshot {
    offered: u64,
    submitted: u64,
    suppressed: u64,
    completed: u64,
    cancelled: u64,
    outstanding: usize,
    latency: HistogramSnapshot,
    task_poll_cpu_seconds: f64,
    task_poll_wall_seconds: f64,
    task_slice_wait_seconds: HistogramSnapshot,
    task_slice_waits: HistogramSnapshot,
}

/// Concurrent simulator-owned collectors for one `(keyspace, task class)`
/// runtime.
#[derive(Clone)]
pub(crate) struct TaskPollMetrics {
    pub(crate) cpu_seconds: Counter,
    pub(crate) wall_seconds: Counter,
    pub(crate) slice_wait_seconds: Histogram,
    pub(crate) slice_waits: Histogram,
}

impl TaskPollMetrics {
    pub(crate) fn new(runtime_id: usize) -> Result<Self> {
        let mut slice_wait_duration_buckets = vec![0.0];
        slice_wait_duration_buckets.extend(exponential_buckets(
            WAIT_BUCKET_START_SECONDS,
            WAIT_BUCKET_FACTOR,
            WAIT_BUCKET_COUNT,
        )?);
        let mut slice_wait_count_buckets = vec![0.0];
        slice_wait_count_buckets.extend(exponential_buckets(
            QUEUE_BUCKET_START,
            QUEUE_BUCKET_FACTOR,
            QUEUE_BUCKET_COUNT,
        )?);
        Ok(Self {
            cpu_seconds: Counter::new(
                format!("yatp_sim_task_poll_cpu_seconds_{runtime_id}_total"),
                "CPU time consumed by YATP task polls for this task runtime",
            )?,
            wall_seconds: Counter::new(
                format!("yatp_sim_task_poll_wall_seconds_{runtime_id}_total"),
                "Wall time consumed by YATP task polls for this task runtime",
            )?,
            slice_wait_seconds: Histogram::with_opts(
                HistogramOpts::new(
                    format!("yatp_sim_task_slice_wait_seconds_{runtime_id}"),
                    "Queue wait before dispatch for this task runtime",
                )
                .buckets(slice_wait_duration_buckets),
            )?,
            slice_waits: Histogram::with_opts(
                HistogramOpts::new(
                    format!("yatp_sim_task_slice_waits_{runtime_id}"),
                    "Successful fair-queue pops while this task runtime waits",
                )
                .buckets(slice_wait_count_buckets),
            )?,
        })
    }
}

struct TaskMetricState {
    offered: u64,
    submitted: u64,
    suppressed: u64,
    completed: u64,
    cancelled: u64,
    outstanding: usize,
    latency: Histogram,
    task_poll_cpu_seconds: Counter,
    task_poll_wall_seconds: Counter,
    task_slice_wait_seconds: Histogram,
    task_slice_waits: Histogram,
}

impl TaskMetricState {
    fn new(runtime_id: usize) -> Result<Self> {
        let latency = Histogram::with_opts(
            HistogramOpts::new(
                format!("yatp_sim_task_latency_seconds_{runtime_id}"),
                "Whole-task latency from successful submission to completion",
            )
            .buckets(TASK_LATENCY_BUCKETS_SECONDS.to_vec()),
        )?;
        let poll_metrics = TaskPollMetrics::new(runtime_id)?;
        Ok(Self {
            offered: 0,
            submitted: 0,
            suppressed: 0,
            completed: 0,
            cancelled: 0,
            outstanding: 0,
            latency,
            task_poll_cpu_seconds: poll_metrics.cpu_seconds,
            task_poll_wall_seconds: poll_metrics.wall_seconds,
            task_slice_wait_seconds: poll_metrics.slice_wait_seconds,
            task_slice_waits: poll_metrics.slice_waits,
        })
    }

    fn snapshot(&self) -> TaskMetricSnapshot {
        TaskMetricSnapshot {
            offered: self.offered,
            submitted: self.submitted,
            suppressed: self.suppressed,
            completed: self.completed,
            cancelled: self.cancelled,
            outstanding: self.outstanding,
            latency: HistogramSnapshot::from_prometheus(&self.latency),
            task_poll_cpu_seconds: self.task_poll_cpu_seconds.get(),
            task_poll_wall_seconds: self.task_poll_wall_seconds.get(),
            task_slice_wait_seconds: HistogramSnapshot::from_prometheus(
                &self.task_slice_wait_seconds,
            ),
            task_slice_waits: HistogramSnapshot::from_prometheus(&self.task_slice_waits),
        }
    }
}

#[derive(Clone, Debug)]
struct CumulativeSnapshot {
    elapsed: Duration,
    worker_cpu_nanos: u64,
    keyspace_queue_lengths: BTreeMap<u32, HistogramSnapshot>,
    tasks: Vec<TaskMetricSnapshot>,
}

#[derive(Clone, Debug)]
struct WorkloadPhaseDefinition {
    active_runtime_ids: Vec<usize>,
    end: Duration,
}

#[derive(Default)]
struct KeyspaceTaskAggregate {
    offered: u64,
    submitted: u64,
    suppressed: u64,
    completed: u64,
    cancelled: u64,
    outstanding: usize,
    latency: Option<HistogramSnapshot>,
    task_poll_cpu_seconds: f64,
    task_poll_wall_seconds: f64,
    task_slice_wait_seconds: Option<HistogramSnapshot>,
    task_slice_waits: Option<HistogramSnapshot>,
}

fn build_workload_phase_definitions(
    config: &SimulatorConfig,
) -> Result<Vec<WorkloadPhaseDefinition>> {
    let mut boundaries = BTreeSet::from([Duration::ZERO, config.duration.0]);
    let mut task_windows = Vec::new();
    let mut runtime_id = 0usize;
    for group in &config.keyspace_groups {
        for _ in 0..group.count {
            for task in &group.tasks {
                if !task.enabled {
                    continue;
                }
                let start = task.start_time.0;
                let end = config.task_end_time(task)?;
                boundaries.insert(start);
                boundaries.insert(end);
                task_windows.push((runtime_id, start, end));
                runtime_id += 1;
            }
        }
    }

    let boundaries = boundaries.into_iter().collect::<Vec<_>>();
    Ok(boundaries
        .windows(2)
        .map(|bounds| WorkloadPhaseDefinition {
            active_runtime_ids: task_windows
                .iter()
                .filter_map(|(runtime_id, start, end)| {
                    (*start < bounds[1] && *end > bounds[0]).then_some(*runtime_id)
                })
                .collect(),
            end: bounds[1],
        })
        .collect())
}

pub struct MetricsRecorder {
    config: SimulatorConfig,
    keyspaces: Vec<KeyspaceDescriptor>,
    tasks: Vec<TaskDescriptor>,
    keyspace_ids: Vec<u32>,
    metric_store: KeyspaceMetricStore,
    worker_cpu_time: Arc<dyn WorkerCpuTimeSource>,
    task_states: Vec<TaskMetricState>,
    initial_snapshot: CumulativeSnapshot,
    last_snapshot: CumulativeSnapshot,
    warmup_snapshot: Option<CumulativeSnapshot>,
    generation_end_snapshot: Option<CumulativeSnapshot>,
    workload_phase_definitions: Vec<WorkloadPhaseDefinition>,
    workload_phase_snapshots: Vec<CumulativeSnapshot>,
    next_workload_phase_boundary: usize,
    windows: Vec<MetricsWindow>,
    next_periodic_at: Duration,
    slow_tasks: Vec<SlowTaskReport>,
}

impl MetricsRecorder {
    pub fn new(
        config: &SimulatorConfig,
        keyspaces: Vec<KeyspaceDescriptor>,
        tasks: Vec<TaskDescriptor>,
        metric_store: KeyspaceMetricStore,
        worker_cpu_time: Arc<dyn WorkerCpuTimeSource>,
    ) -> Result<Self> {
        for (runtime_id, task) in tasks.iter().enumerate() {
            ensure!(
                task.runtime_id == runtime_id,
                "task descriptors must use dense runtime IDs"
            );
        }
        let keyspace_ids = keyspaces
            .iter()
            .map(|keyspace| keyspace.keyspace_id)
            .collect::<Vec<_>>();
        let task_states = tasks
            .iter()
            .map(|task| TaskMetricState::new(task.runtime_id))
            .collect::<Result<Vec<_>>>()?;
        let initial_snapshot = CumulativeSnapshot {
            elapsed: Duration::ZERO,
            worker_cpu_nanos: worker_cpu_time.snapshot_cpu_nanos()?,
            keyspace_queue_lengths: metric_store
                .snapshot_queue_lengths(&keyspace_ids, config.read_pool.keyspace_fairness.enabled),
            tasks: task_states.iter().map(TaskMetricState::snapshot).collect(),
        };
        let warmup_snapshot =
            (config.warmup_duration.0 == Duration::ZERO).then(|| initial_snapshot.clone());
        let workload_phase_definitions = build_workload_phase_definitions(config)?;
        Ok(Self {
            config: config.clone(),
            keyspaces,
            tasks,
            keyspace_ids,
            metric_store,
            worker_cpu_time,
            task_states,
            initial_snapshot: initial_snapshot.clone(),
            last_snapshot: initial_snapshot.clone(),
            warmup_snapshot,
            generation_end_snapshot: None,
            workload_phase_definitions,
            workload_phase_snapshots: vec![initial_snapshot],
            next_workload_phase_boundary: 0,
            windows: Vec::new(),
            next_periodic_at: config.stats_print_interval.0,
            slow_tasks: Vec::new(),
        })
    }

    pub fn record_offered(&mut self, runtime_id: usize, count: u64) {
        self.task_states[runtime_id].offered =
            self.task_states[runtime_id].offered.saturating_add(count);
    }

    pub fn record_suppressed(&mut self, runtime_id: usize, count: u64) {
        self.task_states[runtime_id].suppressed = self.task_states[runtime_id]
            .suppressed
            .saturating_add(count);
    }

    pub fn record_submitted(&mut self, runtime_id: usize) {
        let state = &mut self.task_states[runtime_id];
        state.submitted = state.submitted.saturating_add(1);
        state.outstanding = state.outstanding.saturating_add(1);
    }

    pub fn record_completion(&mut self, runtime_id: usize, completed: bool, latency: Duration) {
        let state = &mut self.task_states[runtime_id];
        debug_assert!(state.outstanding > 0);
        state.outstanding = state.outstanding.saturating_sub(1);
        if completed {
            state.completed = state.completed.saturating_add(1);
            state.latency.observe(latency.as_secs_f64());
        } else {
            state.cancelled = state.cancelled.saturating_add(1);
        }
    }

    pub fn record_slow_task(&mut self, slow_task: SlowTaskReport) {
        self.slow_tasks.push(slow_task);
    }

    /// Returns the simulator-owned poll metrics for one task runtime.
    pub fn task_poll_metrics(&self, runtime_id: usize) -> TaskPollMetrics {
        let state = &self.task_states[runtime_id];
        TaskPollMetrics {
            cpu_seconds: state.task_poll_cpu_seconds.clone(),
            wall_seconds: state.task_poll_wall_seconds.clone(),
            slice_wait_seconds: state.task_slice_wait_seconds.clone(),
            slice_waits: state.task_slice_waits.clone(),
        }
    }

    pub fn next_capture_deadline(&self) -> Duration {
        let mut deadline = self.next_periodic_at;
        if self.warmup_snapshot.is_none() {
            deadline = deadline.min(self.config.warmup_duration.0);
        }
        if self.generation_end_snapshot.is_none() {
            deadline = deadline.min(self.config.duration.0);
        }
        if let Some(phase) = self
            .workload_phase_definitions
            .get(self.next_workload_phase_boundary)
        {
            deadline = deadline.min(phase.end);
        }
        deadline
    }

    pub fn capture_if_due(&mut self, elapsed: Duration) -> Result<()> {
        if elapsed >= self.next_capture_deadline() {
            self.capture(elapsed)?;
        }
        Ok(())
    }

    pub fn capture_generation_end(&mut self, elapsed: Duration) -> Result<()> {
        ensure!(
            self.warmup_snapshot.is_some(),
            "warmup boundary was not captured before workload generation ended"
        );
        self.capture(elapsed.max(self.config.duration.0))?;
        let warmup_elapsed = self
            .warmup_snapshot
            .as_ref()
            .context("warmup snapshot was not captured")?
            .elapsed;
        let generation_end_elapsed = self
            .generation_end_snapshot
            .as_ref()
            .context("generation-end snapshot was not captured")?
            .elapsed;
        ensure!(
            generation_end_elapsed > warmup_elapsed,
            "measurement interval is empty"
        );
        Ok(())
    }

    pub fn finish(mut self, elapsed: Duration) -> Result<ReportData> {
        if self.generation_end_snapshot.is_none() {
            self.capture_generation_end(elapsed)?;
        } else if elapsed > self.last_snapshot.elapsed {
            self.capture(elapsed)?;
        }

        let final_snapshot = self.last_snapshot.clone();
        let warmup_snapshot = self
            .warmup_snapshot
            .clone()
            .context("warmup snapshot was not captured")?;
        let generation_end_snapshot = self
            .generation_end_snapshot
            .clone()
            .context("generation-end snapshot was not captured")?;

        let warmup = (self.config.warmup_duration.0 > Duration::ZERO)
            .then(|| {
                self.build_window(
                    &self.initial_snapshot,
                    &warmup_snapshot,
                    ReportPhase::Warmup,
                )
            })
            .transpose()?;
        let measurement = self.build_window(
            &warmup_snapshot,
            &generation_end_snapshot,
            ReportPhase::Measurement,
        )?;
        let drain = (final_snapshot.elapsed > generation_end_snapshot.elapsed)
            .then(|| {
                self.build_window(
                    &generation_end_snapshot,
                    &final_snapshot,
                    ReportPhase::Drain,
                )
            })
            .transpose()?;
        let total =
            self.build_window(&self.initial_snapshot, &final_snapshot, ReportPhase::Total)?;
        ensure!(
            self.workload_phase_snapshots.len() == self.workload_phase_definitions.len() + 1,
            "workload phase boundaries were not fully captured"
        );
        let workload_phases = self
            .workload_phase_definitions
            .iter()
            .zip(self.workload_phase_snapshots.windows(2))
            .enumerate()
            .map(|(index, (definition, snapshots))| {
                Ok(WorkloadPhaseReport {
                    number: index + 1,
                    active_runtime_ids: definition.active_runtime_ids.clone(),
                    metrics: self.build_window(
                        &snapshots[0],
                        &snapshots[1],
                        ReportPhase::Workload,
                    )?,
                })
            })
            .collect::<Result<Vec<_>>>()?;

        self.slow_tasks.sort_by(|left, right| {
            left.submitted_at_seconds
                .total_cmp(&right.submitted_at_seconds)
                .then_with(|| left.task_id.cmp(&right.task_id))
        });

        Ok(ReportData {
            schema_version: 1,
            generated_at_unix_seconds: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            config: self.config,
            keyspaces: self.keyspaces,
            task_runtimes: self.tasks,
            windows: self.windows,
            summary: RunSummary {
                warmup,
                measurement,
                drain,
                total,
                workload_phases,
            },
            slow_tasks: self.slow_tasks,
        })
    }

    fn capture(&mut self, elapsed: Duration) -> Result<()> {
        if elapsed <= self.last_snapshot.elapsed {
            return Ok(());
        }
        let current = self.snapshot(elapsed)?;
        let phase = self.phase_for(self.last_snapshot.elapsed);
        let window = self.build_window(&self.last_snapshot, &current, phase)?;
        let active_runtime_ids = self
            .workload_phase_definitions
            .get(self.next_workload_phase_boundary)
            .map(|phase| phase.active_runtime_ids.as_slice())
            .unwrap_or_default();
        print_window_summary(&window, active_runtime_ids);
        self.windows.push(window);

        while self
            .workload_phase_definitions
            .get(self.next_workload_phase_boundary)
            .is_some_and(|phase| elapsed >= phase.end)
        {
            self.workload_phase_snapshots.push(current.clone());
            self.next_workload_phase_boundary += 1;
        }

        if self.warmup_snapshot.is_none() && elapsed >= self.config.warmup_duration.0 {
            self.warmup_snapshot = Some(current.clone());
        }
        if self.generation_end_snapshot.is_none() && elapsed >= self.config.duration.0 {
            self.generation_end_snapshot = Some(current.clone());
        }
        self.last_snapshot = current;
        while self.next_periodic_at <= elapsed {
            self.next_periodic_at = self
                .next_periodic_at
                .checked_add(self.config.stats_print_interval.0)
                .unwrap_or(Duration::MAX);
        }
        Ok(())
    }

    fn snapshot(&self, elapsed: Duration) -> Result<CumulativeSnapshot> {
        Ok(CumulativeSnapshot {
            elapsed,
            worker_cpu_nanos: self.worker_cpu_time.snapshot_cpu_nanos()?,
            keyspace_queue_lengths: self.metric_store.snapshot_queue_lengths(
                &self.keyspace_ids,
                self.config.read_pool.keyspace_fairness.enabled,
            ),
            tasks: self
                .task_states
                .iter()
                .map(TaskMetricState::snapshot)
                .collect(),
        })
    }

    fn phase_for(&self, start: Duration) -> ReportPhase {
        if start < self.config.warmup_duration.0 {
            ReportPhase::Warmup
        } else if start < self.config.duration.0 {
            ReportPhase::Measurement
        } else {
            ReportPhase::Drain
        }
    }

    fn build_window(
        &self,
        before: &CumulativeSnapshot,
        after: &CumulativeSnapshot,
        phase: ReportPhase,
    ) -> Result<MetricsWindow> {
        ensure!(
            after.elapsed >= before.elapsed,
            "metrics window runs backwards"
        );
        ensure!(
            before.tasks.len() == after.tasks.len() && after.tasks.len() == self.tasks.len(),
            "task metrics snapshot shape changed"
        );
        let duration = after.elapsed.saturating_sub(before.elapsed);
        let duration_seconds = duration.as_secs_f64().max(f64::EPSILON);
        let total_cpu_seconds = after
            .worker_cpu_nanos
            .saturating_sub(before.worker_cpu_nanos) as f64
            / 1_000_000_000.0;
        let mut keyspace_aggregates = self
            .keyspaces
            .iter()
            .map(|keyspace| (keyspace.keyspace_id, KeyspaceTaskAggregate::default()))
            .collect::<BTreeMap<_, _>>();

        let mut tasks = Vec::with_capacity(self.tasks.len());
        for ((descriptor, before), after) in self.tasks.iter().zip(&before.tasks).zip(&after.tasks)
        {
            let latency = after.latency.delta(&before.latency);
            let offered = after.offered.saturating_sub(before.offered);
            let submitted = after.submitted.saturating_sub(before.submitted);
            let suppressed = after.suppressed.saturating_sub(before.suppressed);
            let completed = after.completed.saturating_sub(before.completed);
            let cancelled = after.cancelled.saturating_sub(before.cancelled);
            let task_poll_cpu_seconds =
                (after.task_poll_cpu_seconds - before.task_poll_cpu_seconds).max(0.0);
            let task_poll_wall_seconds =
                (after.task_poll_wall_seconds - before.task_poll_wall_seconds).max(0.0);
            let task_slice_wait_seconds = after
                .task_slice_wait_seconds
                .delta(&before.task_slice_wait_seconds);
            let task_slice_waits = after.task_slice_waits.delta(&before.task_slice_waits);
            let aggregate = keyspace_aggregates
                .get_mut(&descriptor.keyspace_id)
                .context("task references an unknown keyspace")?;
            aggregate.offered = aggregate.offered.saturating_add(offered);
            aggregate.submitted = aggregate.submitted.saturating_add(submitted);
            aggregate.suppressed = aggregate.suppressed.saturating_add(suppressed);
            aggregate.completed = aggregate.completed.saturating_add(completed);
            aggregate.cancelled = aggregate.cancelled.saturating_add(cancelled);
            aggregate.outstanding = aggregate.outstanding.saturating_add(after.outstanding);
            match aggregate.latency.as_mut() {
                Some(existing) => existing.merge(&latency),
                None => aggregate.latency = Some(latency.clone()),
            }
            aggregate.task_poll_cpu_seconds += task_poll_cpu_seconds;
            aggregate.task_poll_wall_seconds += task_poll_wall_seconds;
            match aggregate.task_slice_wait_seconds.as_mut() {
                Some(existing) => existing.merge(&task_slice_wait_seconds),
                None => aggregate.task_slice_wait_seconds = Some(task_slice_wait_seconds.clone()),
            }
            match aggregate.task_slice_waits.as_mut() {
                Some(existing) => existing.merge(&task_slice_waits),
                None => aggregate.task_slice_waits = Some(task_slice_waits.clone()),
            }
            tasks.push(TaskMetricsReport {
                runtime_id: descriptor.runtime_id,
                keyspace_id: descriptor.keyspace_id,
                keyspace_name: descriptor.keyspace_name.clone(),
                group_name: descriptor.group_name.clone(),
                group_instance: descriptor.group_instance,
                task_name: descriptor.task_name.clone(),
                offered,
                submitted,
                suppressed,
                completed,
                cancelled,
                outstanding_at_end: after.outstanding,
                throughput_qps: completed as f64 / duration_seconds,
                task_latency_seconds: latency.report(),
                task_poll_cpu_seconds,
                task_poll_wall_seconds,
                task_poll_cpu_cores: task_poll_cpu_seconds / duration_seconds,
                task_poll_wall_cores: task_poll_wall_seconds / duration_seconds,
                task_slice_wait_seconds: task_slice_wait_seconds.report(),
                task_slice_waits: task_slice_waits.report(),
            });
        }

        let mut keyspaces = Vec::with_capacity(self.keyspaces.len());
        for descriptor in &self.keyspaces {
            let aggregate = keyspace_aggregates
                .remove(&descriptor.keyspace_id)
                .context("missing keyspace task aggregate")?;
            let fair_queue_length = match (
                before.keyspace_queue_lengths.get(&descriptor.keyspace_id),
                after.keyspace_queue_lengths.get(&descriptor.keyspace_id),
            ) {
                (Some(before), Some(after)) => Some(after.delta(before).report()),
                (None, None) => None,
                _ => anyhow::bail!(
                    "keyspace queue-length metric availability changed during the run"
                ),
            };
            let cpu_seconds = aggregate.task_poll_cpu_seconds;
            let wall_seconds = aggregate.task_poll_wall_seconds;
            let accounted_seconds = self.config.read_pool.keyspace_fairness.enabled.then_some(
                if self.config.read_pool.keyspace_fairness.use_wall_time {
                    wall_seconds
                } else {
                    cpu_seconds
                },
            );
            let (task_slice_waits, task_slice_wait_seconds) =
                if self.config.read_pool.keyspace_fairness.enabled {
                    (
                        aggregate.task_slice_waits.map(|metrics| metrics.report()),
                        aggregate
                            .task_slice_wait_seconds
                            .map(|metrics| metrics.report()),
                    )
                } else {
                    (None, None)
                };
            keyspaces.push(KeyspaceMetricsReport {
                keyspace_id: descriptor.keyspace_id,
                keyspace_name: descriptor.keyspace_name.clone(),
                group_name: descriptor.group_name.clone(),
                group_instance: descriptor.group_instance,
                offered: aggregate.offered,
                submitted: aggregate.submitted,
                suppressed: aggregate.suppressed,
                completed: aggregate.completed,
                cancelled: aggregate.cancelled,
                outstanding_at_end: aggregate.outstanding,
                throughput_qps: aggregate.completed as f64 / duration_seconds,
                task_latency_seconds: aggregate
                    .latency
                    .unwrap_or_else(HistogramSnapshot::empty)
                    .report(),
                task_poll_cpu_seconds: cpu_seconds,
                task_poll_wall_seconds: wall_seconds,
                task_poll_cpu_cores: cpu_seconds / duration_seconds,
                task_poll_wall_cores: wall_seconds / duration_seconds,
                accounted_seconds,
                accounted_service_share: None,
                fair_queue_length,
                fair_queue_task_slice_waits: task_slice_waits,
                fair_queue_task_slice_wait_seconds: task_slice_wait_seconds,
            });
        }

        let total_accounted_seconds = keyspaces
            .iter()
            .filter_map(|keyspace| keyspace.accounted_seconds)
            .sum::<f64>();
        if self.config.read_pool.keyspace_fairness.enabled {
            for keyspace in &mut keyspaces {
                keyspace.accounted_service_share = Some(if total_accounted_seconds > 0.0 {
                    keyspace.accounted_seconds.unwrap_or(0.0) / total_accounted_seconds
                } else {
                    0.0
                });
            }
        }

        Ok(MetricsWindow {
            phase,
            start_seconds: before.elapsed.as_secs_f64(),
            end_seconds: after.elapsed.as_secs_f64(),
            duration_seconds,
            total_cpu_seconds,
            total_cpu_cores: total_cpu_seconds / duration_seconds,
            keyspaces,
            tasks,
        })
    }
}

fn print_window_summary(window: &MetricsWindow, active_runtime_ids: &[usize]) {
    if active_runtime_ids.is_empty() {
        return;
    }
    println!(
        "{}",
        format_window_summary_for_runtime_ids(window, Some(active_runtime_ids))
    );
}

struct TaskSummaryLine {
    workload: String,
    task_lat_avg: String,
    task_lat_p99: String,
    cpu_cores: String,
    completed: String,
    qps: String,
    suppressed: String,
    cancelled: String,
}

impl TaskSummaryLine {
    fn from_report(task: &TaskMetricsReport) -> Self {
        Self {
            workload: format!("[{}/{}]", task.keyspace_name, task.task_name),
            task_lat_avg: format_latency(task.task_latency_seconds.average),
            task_lat_p99: format_latency(task.task_latency_seconds.p99),
            cpu_cores: format!("{:.3}", task.task_poll_cpu_cores),
            completed: task.completed.to_string(),
            qps: format!("{:.3}", task.throughput_qps),
            suppressed: task.suppressed.to_string(),
            cancelled: task.cancelled.to_string(),
        }
    }
}

fn format_latency(seconds: Option<f64>) -> String {
    let Some(seconds) = seconds else {
        return "N/A".to_owned();
    };
    if seconds < 0.001 {
        format!("{}us", format_compact_decimal(seconds * 1_000_000.0, 3))
    } else if seconds < 1.0 {
        format!("{}ms", format_compact_decimal(seconds * 1_000.0, 3))
    } else {
        format!("{}s", format_compact_decimal(seconds, 3))
    }
}

fn format_compact_decimal(value: f64, decimal_places: usize) -> String {
    let mut formatted = format!("{value:.decimal_places$}");
    if formatted.contains('.') {
        while formatted.ends_with('0') {
            formatted.pop();
        }
        if formatted.ends_with('.') {
            formatted.pop();
        }
    }
    if formatted == "-0" {
        "0".to_owned()
    } else {
        formatted
    }
}

fn format_interval_timestamp(seconds: f64) -> String {
    if seconds < 0.001 {
        format!("{}us", format_compact_decimal(seconds * 1_000_000.0, 0))
    } else if seconds < 1.0 {
        format!("{}ms", format_compact_decimal(seconds * 1_000.0, 0))
    } else {
        format!("{}s", format_compact_decimal(seconds, 0))
    }
}

#[cfg(test)]
fn format_window_summary(window: &MetricsWindow) -> String {
    format_window_summary_for_runtime_ids(window, None)
}

fn format_window_summary_for_runtime_ids(
    window: &MetricsWindow,
    active_runtime_ids: Option<&[usize]>,
) -> String {
    let mut output = format!("[{}]", format_interval_timestamp(window.end_seconds));
    if window.phase == ReportPhase::Warmup {
        output.push_str(" (WARMUP)");
    }
    write!(
        &mut output,
        " total_cpu_cores={}",
        format_compact_decimal(window.total_cpu_cores, 3),
    )
    .expect("writing an interval summary to a String cannot fail");
    let rows = window
        .tasks
        .iter()
        .filter(|task| {
            active_runtime_ids
                .map(|runtime_ids| runtime_ids.contains(&task.runtime_id))
                .unwrap_or(true)
        })
        .map(TaskSummaryLine::from_report)
        .collect::<Vec<_>>();
    let workload_width = rows.iter().map(|row| row.workload.len()).max().unwrap_or(0);
    let task_lat_avg_width = rows
        .iter()
        .map(|row| row.task_lat_avg.len())
        .max()
        .unwrap_or(0);
    let task_lat_p99_width = rows
        .iter()
        .map(|row| row.task_lat_p99.len())
        .max()
        .unwrap_or(0);
    let cpu_cores_width = rows
        .iter()
        .map(|row| row.cpu_cores.len())
        .max()
        .unwrap_or(0);
    let completed_width = rows
        .iter()
        .map(|row| row.completed.len())
        .max()
        .unwrap_or(0);
    let qps_width = rows.iter().map(|row| row.qps.len()).max().unwrap_or(0);
    let suppressed_width = rows
        .iter()
        .map(|row| row.suppressed.len())
        .max()
        .unwrap_or(0);

    for row in rows {
        write!(
            &mut output,
            "\n    {:<workload_width$} task_lat_avg={:>task_lat_avg_width$} task_lat_p99={:>task_lat_p99_width$} cpu_cores={:>cpu_cores_width$} completed={:>completed_width$} qps={:>qps_width$} suppressed={:>suppressed_width$} cancelled={}",
            row.workload,
            row.task_lat_avg,
            row.task_lat_p99,
            row.cpu_cores,
            row.completed,
            row.qps,
            row.suppressed,
            row.cancelled,
        )
        .expect("writing a terminal summary to a String cannot fail");
    }
    output
}

/// Prints aggregate statistics for each configured workload combination.
pub fn print_run_summary(report: &ReportData) {
    println!("{}", format_run_summary(report));
}

fn format_run_summary(report: &ReportData) -> String {
    let mut output = "SUMMARY".to_owned();
    for phase in &report.summary.workload_phases {
        write_workload_phase_summary(&mut output, phase, report.config.read_pool.workers);
    }
    write!(&mut output, "\n  Slow tasks: {}", report.slow_tasks.len())
        .expect("writing a slow-task count to a String cannot fail");
    output
}

fn write_workload_phase_summary(output: &mut String, phase: &WorkloadPhaseReport, workers: usize) {
    let window = &phase.metrics;
    let tasks = phase
        .active_runtime_ids
        .iter()
        .filter_map(|runtime_id| window.tasks.get(*runtime_id))
        .collect::<Vec<_>>();
    let workload_names = tasks
        .iter()
        .map(|task| task_runtime_label(task))
        .collect::<Vec<_>>()
        .join(", ");
    write!(
        output,
        "\n  Phase{}: [{}], interval=[{:.3}s, {:.3}s) duration={:.3}s",
        phase.number,
        workload_names,
        window.start_seconds,
        window.end_seconds,
        window.duration_seconds,
    )
    .expect("writing a terminal summary to a String cannot fail");

    let worker_capacity = workers as f64;
    write!(
        output,
        "\n    Total CPU: {} cores, utilization: {}%",
        format_compact_decimal(window.total_cpu_cores, 3),
        format_compact_decimal(window.total_cpu_cores / worker_capacity * 100.0, 1),
    )
    .expect("writing a terminal summary to a String cannot fail");

    for task in &tasks {
        write!(
            output,
            "\n    [{}] QPS: {}, CPU: {} cores, utilization: {}%",
            task_runtime_label(task),
            format_compact_decimal(task.throughput_qps, 3),
            format_compact_decimal(task.task_poll_cpu_cores, 3),
            format_compact_decimal(task.task_poll_cpu_cores / worker_capacity * 100.0, 1),
        )
        .expect("writing a terminal summary to a String cannot fail");
        write!(
            output,
            "\n          Latency: avg={}, p50={}, p80={}, p95={}, p99={}",
            format_latency(task.task_latency_seconds.average),
            format_latency(task.task_latency_seconds.p50),
            format_latency(task.task_latency_seconds.p80),
            format_latency(task.task_latency_seconds.p95),
            format_latency(task.task_latency_seconds.p99),
        )
        .expect("writing a terminal summary to a String cannot fail");
    }
}

fn task_runtime_label(task: &TaskMetricsReport) -> String {
    format!("{}/{}", task.keyspace_name, task.task_name)
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicU64, Ordering};

    use tikv_util::config::ReadableDuration;

    use super::*;

    #[derive(Default)]
    struct TestWorkerCpuTimeSource {
        cpu_nanos: AtomicU64,
    }

    impl WorkerCpuTimeSource for TestWorkerCpuTimeSource {
        fn snapshot_cpu_nanos(&self) -> Result<u64> {
            Ok(self.cpu_nanos.load(Ordering::Relaxed))
        }
    }

    fn test_config() -> SimulatorConfig {
        toml::from_str(
            r#"
duration = "1s"
stats-print-interval = "100ms"

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

[[keyspace-groups.tasks]]
name = "read"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5
"#,
        )
        .unwrap()
    }

    fn descriptors() -> (Vec<KeyspaceDescriptor>, Vec<TaskDescriptor>) {
        (
            vec![KeyspaceDescriptor {
                keyspace_id: 1,
                keyspace_name: "test".to_owned(),
                group_name: "test".to_owned(),
                group_instance: 1,
            }],
            vec![TaskDescriptor {
                runtime_id: 0,
                keyspace_id: 1,
                keyspace_name: "test".to_owned(),
                group_name: "test".to_owned(),
                group_instance: 1,
                task_name: "read".to_owned(),
            }],
        )
    }

    fn test_metrics_recorder(
        config: &SimulatorConfig,
        keyspaces: Vec<KeyspaceDescriptor>,
        tasks: Vec<TaskDescriptor>,
    ) -> MetricsRecorder {
        MetricsRecorder::new(
            config,
            keyspaces,
            tasks,
            KeyspaceMetricStore::new().unwrap(),
            Arc::new(TestWorkerCpuTimeSource::default()),
        )
        .unwrap()
    }

    #[test]
    fn test_production_keyspace_queue_length_is_reported() {
        let mut config = test_config();
        config.read_pool.keyspace_fairness.enabled = true;
        let (keyspaces, tasks) = descriptors();
        let metric_store = KeyspaceMetricStore::new().unwrap();
        let keyspace_metrics = metric_store.keyspace_metrics().get_keyspace_metrics(1);
        let mut recorder = MetricsRecorder::new(
            &config,
            keyspaces,
            tasks,
            metric_store,
            Arc::new(TestWorkerCpuTimeSource::default()),
        )
        .unwrap();

        keyspace_metrics.observe_fair_queue_length(3);
        recorder
            .capture_if_due(ReadableDuration::millis(100).0)
            .unwrap();
        keyspace_metrics.observe_fair_queue_length(7);
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();
        let queue_length = report.summary.measurement.keyspaces[0]
            .fair_queue_length
            .as_ref()
            .unwrap();
        assert_eq!(queue_length.count, 2);
        assert_eq!(queue_length.p50, Some(4.0));
        assert_eq!(queue_length.p99, Some(8.0));
        assert_eq!(
            report.windows[0].keyspaces[0]
                .fair_queue_length
                .as_ref()
                .unwrap()
                .p99,
            Some(4.0)
        );
        assert_eq!(
            report.windows[1].keyspaces[0]
                .fair_queue_length
                .as_ref()
                .unwrap()
                .p99,
            Some(8.0)
        );
    }

    #[test]
    fn test_cancelled_task_is_excluded_from_latency() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);

        recorder.record_offered(0, 2);
        recorder.record_submitted(0);
        recorder.record_submitted(0);
        recorder.record_completion(0, true, Duration::from_millis(10));
        recorder.record_completion(0, false, Duration::from_millis(1));
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();
        let task = &report.summary.measurement.tasks[0];
        assert_eq!(task.offered, 2);
        assert_eq!(task.submitted, 2);
        assert_eq!(task.completed, 1);
        assert_eq!(task.cancelled, 1);
        assert_eq!(task.task_latency_seconds.count, 1);
        assert!(
            report.summary.measurement.keyspaces[0]
                .fair_queue_length
                .is_none()
        );
        assert!(report.summary.drain.is_none());
    }

    #[test]
    fn test_task_poll_metrics_delta_is_reported() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);
        let task_poll = recorder.task_poll_metrics(0);

        task_poll.cpu_seconds.inc_by(0.010);
        task_poll.wall_seconds.inc_by(0.012);
        task_poll.slice_wait_seconds.observe(0.002);
        task_poll.slice_waits.observe(2.0);
        recorder
            .capture_if_due(ReadableDuration::millis(100).0)
            .unwrap();
        task_poll.cpu_seconds.inc_by(0.025);
        task_poll.wall_seconds.inc_by(0.030);
        task_poll.slice_wait_seconds.observe(0.008);
        task_poll.slice_waits.observe(8.0);
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();
        let task = &report.summary.measurement.tasks[0];
        assert!((task.task_poll_cpu_seconds - 0.035).abs() < 1e-12);
        assert!((task.task_poll_cpu_cores - 0.035).abs() < 1e-12);
        assert!((task.task_poll_wall_seconds - 0.042).abs() < 1e-12);
        assert!((task.task_poll_wall_cores - 0.042).abs() < 1e-12);
        assert_eq!(task.task_slice_wait_seconds.count, 2);
        assert_eq!(task.task_slice_waits.count, 2);

        let second_window_task = &report.windows[1].tasks[0];
        assert!((second_window_task.task_poll_cpu_seconds - 0.025).abs() < 1e-12);
        assert!((second_window_task.task_poll_cpu_cores - 0.025 / 0.9).abs() < 1e-12);
        assert!((second_window_task.task_poll_wall_seconds - 0.030).abs() < 1e-12);
        assert_eq!(second_window_task.task_slice_wait_seconds.count, 1);
        assert_eq!(second_window_task.task_slice_waits.count, 1);
    }

    #[test]
    fn test_total_cpu_uses_worker_threads_instead_of_task_sum() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let worker_cpu_time = Arc::new(TestWorkerCpuTimeSource::default());
        let mut recorder = MetricsRecorder::new(
            &config,
            keyspaces,
            tasks,
            KeyspaceMetricStore::new().unwrap(),
            worker_cpu_time.clone(),
        )
        .unwrap();
        recorder.task_poll_metrics(0).cpu_seconds.inc_by(0.010);
        worker_cpu_time
            .cpu_nanos
            .store(50_000_000, Ordering::Relaxed);

        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();
        let measurement = &report.summary.measurement;
        assert!((measurement.total_cpu_seconds - 0.050).abs() < 1e-12);
        assert!((measurement.total_cpu_cores - 0.050).abs() < 1e-12);
        assert!((measurement.tasks[0].task_poll_cpu_seconds - 0.010).abs() < 1e-12);
    }

    #[test]
    fn test_task_poll_wait_histograms_preserve_zero() {
        let metrics = TaskPollMetrics::new(100).unwrap();
        metrics.slice_wait_seconds.observe(0.0);
        metrics.slice_waits.observe(0.0);

        let wait_seconds = HistogramSnapshot::from_prometheus(&metrics.slice_wait_seconds).report();
        let waits = HistogramSnapshot::from_prometheus(&metrics.slice_waits).report();
        assert_eq!(wait_seconds.p99, Some(0.0));
        assert_eq!(waits.p99, Some(0.0));
    }

    #[test]
    fn test_window_summary_groups_tasks_under_one_header() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();

        let summary = format_window_summary(&report.summary.measurement);
        let lines = summary.lines().collect::<Vec<_>>();
        assert_eq!(lines.len(), 2);
        assert_eq!(lines[0], "[1s] total_cpu_cores=0");
        assert!(lines[1].starts_with("    [test/read] task_lat_avg="));
        assert!(lines[1].contains("task_lat_avg=N/A task_lat_p99=N/A"));
        assert!(!lines[1].contains("phase="));

        let mut warmup = report.summary.measurement;
        warmup.phase = ReportPhase::Warmup;
        warmup.end_seconds = 5.001;
        warmup.total_cpu_cores = 4.0;
        let warmup_summary = format_window_summary(&warmup);
        assert_eq!(
            warmup_summary.lines().next().unwrap(),
            "[5s] (WARMUP) total_cpu_cores=4"
        );
    }

    #[test]
    fn test_window_summary_aligns_task_columns() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);
        recorder.record_submitted(0);
        recorder.record_completion(0, true, Duration::from_millis(20));
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();
        let mut window = report.summary.measurement;
        let mut second = window.tasks[0].clone();
        second.runtime_id = 1;
        second.keyspace_name = "long-keyspace".to_owned();
        second.task_name = "long-task-name".to_owned();
        second.task_latency_seconds.p99 = Some(1.0);
        second.task_poll_cpu_cores = 12.0;
        second.completed = 1_000;
        second.throughput_qps = 1_000.0;
        window.tasks.push(second);

        let summary = format_window_summary(&window);
        let lines = summary.lines().skip(1).collect::<Vec<_>>();
        assert_eq!(lines.len(), 2);
        for key in [
            "task_lat_avg=",
            "task_lat_p99=",
            "cpu_cores=",
            "completed=",
            "qps=",
            "suppressed=",
            "cancelled=",
        ] {
            assert_eq!(lines[0].find(key), lines[1].find(key), "unaligned {key}");
        }
        assert!(lines[0].contains("task_lat_avg=20ms task_lat_p99=20ms"));
        assert!(lines[1].contains("task_lat_avg=20ms task_lat_p99=  1s"));

        let filtered = format_window_summary_for_runtime_ids(&window, Some(&[0]));
        assert!(filtered.contains("[test/read]"));
        assert!(!filtered.contains("[long-keyspace/long-task-name]"));
    }

    #[test]
    fn test_report_timeline_excludes_cooldown() {
        let config = test_config();
        assert_eq!(config.cooldown_duration, ReadableDuration::secs(5));
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);

        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();

        assert_eq!(report.summary.total.end_seconds, 1.0);
        assert_eq!(report.windows.last().unwrap().end_seconds, 1.0);
        assert!(report.summary.drain.is_none());
    }

    #[test]
    fn test_quantile_is_unavailable_when_it_falls_in_overflow() {
        let histogram = HistogramSnapshot {
            count: 2,
            sum: 3.0,
            buckets: vec![HistogramBucketSnapshot {
                upper_bound: 1.0,
                cumulative_count: 1,
            }],
        };

        let report = histogram.report();
        assert_eq!(report.overflow_count, 1);
        assert_eq!(report.p50, Some(1.0));
        assert_eq!(report.p80, None);
        assert_eq!(report.p95, None);
        assert_eq!(report.p99, None);
    }

    #[test]
    fn test_task_latency_p99_uses_fine_grained_buckets() {
        let state = TaskMetricState::new(42).unwrap();
        for _ in 0..98 {
            state.latency.observe(0.005);
        }
        for _ in 0..2 {
            state.latency.observe(0.012);
        }

        let report = state.snapshot().latency.report();
        let p99 = report.p99.unwrap();
        assert_eq!(p99, 0.015);
        assert_eq!(report.buckets.last().unwrap().upper_bound, 600.0);
    }

    #[test]
    fn test_run_summary_reports_each_task_and_total_cpu() {
        let config = test_config();
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);
        recorder.record_offered(0, 1);
        recorder.record_submitted(0);
        recorder.record_completion(0, true, Duration::from_millis(12));
        recorder.record_slow_task(SlowTaskReport {
            task_id: 1,
            runtime_id: 0,
            offer_sequence: 0,
            submitted_at_seconds: 0.0,
            completed_at_seconds: 0.012,
            latency_seconds: 0.012,
            slow_threshold_seconds: 0.01,
            polls: Vec::new(),
        });
        recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap();
        let report = recorder.finish(ReadableDuration::secs(1).0).unwrap();

        let output = format_run_summary(&report);
        assert!(output.starts_with(
            "SUMMARY\n  Phase1: [test/read], interval=[0.000s, 1.000s) duration=1.000s"
        ));
        assert!(output.contains("[test/read] QPS: 1, CPU: 0 cores, utilization: 0%"));
        assert!(output.contains("          Latency: avg=12ms, p50=15ms, p80=15ms"));
        assert!(output.contains("Total CPU: 0 cores, utilization: 0%"));
        assert!(output.ends_with("  Slow tasks: 1"));
        assert!(
            output.find("Total CPU:").unwrap() < output.find("[test/read] QPS:").unwrap(),
            "total CPU should be printed before workload details"
        );
        assert!(!output.contains("phase=measurement"));
    }

    #[test]
    fn test_compact_decimal_trims_only_insignificant_digits() {
        assert_eq!(format_compact_decimal(199.940, 3), "199.94");
        assert_eq!(format_compact_decimal(1.027, 3), "1.027");
        assert_eq!(format_compact_decimal(10.0, 3), "10");
        assert_eq!(format_compact_decimal(-0.0001, 3), "0");
        assert_eq!(format_interval_timestamp(0.4), "400ms");
        assert_eq!(format_interval_timestamp(0.0004), "400us");
    }

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

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

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

[[keyspace-groups.tasks]]
name = "burst"
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5
start-time = "400ms"
duration = "300ms"

[[keyspace-groups.tasks]]
name = "disabled"
enabled = false
arrival-mode = "closed-loop"
concurrency = 1
task-cpu-us = 10
slice-cpu-us = 5
start-time = "200ms"
duration = "100ms"
"#,
        )
        .unwrap();

        let phases = build_workload_phase_definitions(&config).unwrap();
        assert_eq!(phases.len(), 3);
        assert_eq!(phases[0].active_runtime_ids, vec![0]);
        assert_eq!(phases[0].end, Duration::from_millis(400));
        assert_eq!(phases[1].active_runtime_ids, vec![0, 1]);
        assert_eq!(phases[1].end, Duration::from_millis(700));
        assert_eq!(phases[2].active_runtime_ids, vec![0]);
        assert_eq!(phases[2].end, Duration::from_secs(1));
    }

    #[test]
    fn test_generation_end_rejects_a_missed_warmup_boundary() {
        let mut config = test_config();
        config.warmup_duration = ReadableDuration::millis(500);
        let (keyspaces, tasks) = descriptors();
        let mut recorder = test_metrics_recorder(&config, keyspaces, tasks);

        let error = recorder
            .capture_generation_end(ReadableDuration::secs(1).0)
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("warmup boundary was not captured"),
            "{error}"
        );
    }
}