// Copyright 2026 PingCAP, Inc.

use std::{ops::Deref, sync::Arc};

use collections::{HashMap, HashSet};
use lazy_static::lazy_static;
use parking_lot::Mutex;
use prometheus::{
    GaugeVec, Opts,
    core::{Collector, Desc},
    proto::MetricFamily,
};
use tikv_util::time::{Duration, Instant};

// The duration that keyspace cpu count not updated will be removed.
const EVICT_STALE_METRIC_DUR: Duration = Duration::from_secs(600);
// Only expose the hottest keyspaces in Prometheus to keep the metrics payload
// bounded.
const REPORTED_KEYSPACE_CPU_TOP_N: usize = 20;
// Report average CPU cores used over a recent sliding window so the value is
// stable across scrape timing. Keep it aligned with Prometheus's default
// 15-second scrape interval.
const RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS: usize = 15;

lazy_static! {
    static ref KEY_SPACE_CPU_COLLECTOR: KeyspaceCpuCollector = KeyspaceCpuCollector::new();
}

pub fn register_keyspace_cpu_metric() {
    if let Err(e) = prometheus::register(Box::new(KEY_SPACE_CPU_COLLECTOR.clone())) {
        tikv_util::warn!("register keyspace cpu metrics failed"; "err" => ?e);
    }
}

pub fn report_keyspace_cpu_time(usages: Vec<(u32, u64)>) {
    KEY_SPACE_CPU_COLLECTOR.record_cpu_time(usages, Instant::now_coarse());
}

#[derive(Clone, Copy, Default)]
struct CpuUsageBucket {
    second: i64,
    cpu_us: u64,
}

struct KeyspaceCpuUsage {
    cpu_usage_buckets: [CpuUsageBucket; RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS],
    updated_at: Instant,
}

impl KeyspaceCpuUsage {
    fn new(now: Instant) -> Self {
        Self {
            cpu_usage_buckets: [CpuUsageBucket::default(); RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS],
            updated_at: now,
        }
    }

    fn add_cpu_time(&mut self, now: Instant, cpu_us: u64) {
        let current_second = now.second();
        let slot = current_second.rem_euclid(RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS as i64) as usize;
        let bucket = &mut self.cpu_usage_buckets[slot];
        if bucket.second != current_second {
            *bucket = CpuUsageBucket {
                second: current_second,
                cpu_us: 0,
            };
        }
        bucket.cpu_us += cpu_us;
        self.updated_at = now;
    }

    fn recent_cpu_us(&self, current_second: i64) -> f64 {
        let total_cpu_us: u64 = self
            .cpu_usage_buckets
            .iter()
            .filter(|bucket| {
                current_second.saturating_sub(bucket.second)
                    < RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS as i64
            })
            .map(|bucket| bucket.cpu_us)
            .sum();
        total_cpu_us as f64 / RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS as f64
    }
}

#[derive(Default)]
struct KeyspaceCpuCollectorState {
    cpu_usage: HashMap<u32, KeyspaceCpuUsage>,
    reported_keyspaces: HashSet<u32>,
}

#[derive(Clone)]
struct KeyspaceCpuCollector {
    inner: Arc<KeyspaceCpuCollectorInner>,
}

impl Deref for KeyspaceCpuCollector {
    type Target = KeyspaceCpuCollectorInner;
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

struct KeyspaceCpuCollectorInner {
    metrics: GaugeVec,
    state: Mutex<KeyspaceCpuCollectorState>,
}

impl KeyspaceCpuCollector {
    fn new() -> Self {
        let metrics = GaugeVec::new(
            Opts::new(
                "tikv_coprocessor_keyspace_handle_cpu_cores",
                "Average coprocessor CPU cores used by keyspace over the recent window.",
            ),
            &["keyspace_id"],
        )
        .unwrap();
        Self {
            inner: Arc::new(KeyspaceCpuCollectorInner {
                metrics,
                state: Mutex::default(),
            }),
        }
    }

    fn cpu_cores(recent_cpu_us: f64) -> f64 {
        recent_cpu_us / 1_000_000.0
    }

    fn record_cpu_time(&self, keyspace_usages: Vec<(u32, u64)>, now: Instant) {
        let mut state = self.state.lock();
        for (id, cpu_us) in keyspace_usages {
            if cpu_us == 0 {
                continue;
            }
            state
                .cpu_usage
                .entry(id)
                .or_insert_with(|| KeyspaceCpuUsage::new(now))
                .add_cpu_time(now, cpu_us);
        }
    }

    fn collect_with_time(&self, now: Instant) -> Vec<MetricFamily> {
        let mut state = self.state.lock();
        state.cpu_usage.retain(|_, stat| {
            now.saturating_duration_since(stat.updated_at) < EVICT_STALE_METRIC_DUR
        });

        let current_second = now.second();
        let mut top_keyspaces: Vec<_> = state
            .cpu_usage
            .iter()
            .filter_map(|(&keyspace_id, stat)| {
                let recent_cpu_us = stat.recent_cpu_us(current_second);
                (recent_cpu_us > 0.0).then_some((keyspace_id, recent_cpu_us))
            })
            .collect();
        top_keyspaces
            .sort_unstable_by(|lhs, rhs| rhs.1.total_cmp(&lhs.1).then_with(|| lhs.0.cmp(&rhs.0)));
        top_keyspaces.truncate(REPORTED_KEYSPACE_CPU_TOP_N);

        let reported_keyspace_ids: HashSet<_> = top_keyspaces
            .iter()
            .map(|(keyspace_id, _)| *keyspace_id)
            .collect();
        for &keyspace_id in state.reported_keyspaces.difference(&reported_keyspace_ids) {
            let id_str = keyspace_id.to_string();
            let _ = self.metrics.remove_label_values(&[&id_str]);
        }
        for (keyspace_id, recent_cpu_us) in top_keyspaces {
            let id_str = keyspace_id.to_string();
            self.metrics
                .get_metric_with_label_values(&[&id_str])
                .unwrap()
                .set(Self::cpu_cores(recent_cpu_us));
        }
        state.reported_keyspaces = reported_keyspace_ids;
        drop(state);

        self.metrics.collect()
    }
}

impl Collector for KeyspaceCpuCollector {
    fn desc(&self) -> Vec<&Desc> {
        self.metrics.desc()
    }

    fn collect(&self) -> Vec<MetricFamily> {
        self.collect_with_time(Instant::now_coarse())
    }
}

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

    fn metric_keyspace_id(metric: &prometheus::proto::Metric) -> u32 {
        let label = &metric.get_label()[0];
        assert_eq!(label.get_name(), "keyspace_id");
        label.get_value().parse().unwrap()
    }

    fn metric_cpu_cores(metric: &prometheus::proto::Metric) -> f64 {
        metric.get_gauge().get_value()
    }

    fn collect_keyspace_cpu(metric_families: Vec<MetricFamily>) -> HashMap<u32, f64> {
        metric_families
            .into_iter()
            .flat_map(|metric_family| {
                metric_family
                    .get_metric()
                    .iter()
                    .map(|metric| (metric_keyspace_id(metric), metric_cpu_cores(metric)))
                    .collect::<Vec<_>>()
            })
            .collect()
    }

    #[test]
    fn test_collect_only_reports_top_n_keyspaces() {
        let collector = KeyspaceCpuCollector::new();
        let now = Instant::from_timespec_second_coarse(100);
        let mut usages = Vec::new();
        for keyspace_id in 0..=(REPORTED_KEYSPACE_CPU_TOP_N as u32) {
            usages.push((keyspace_id, (keyspace_id + 1) as u64));
        }
        collector.record_cpu_time(usages, now);

        let keyspace_cpu = collect_keyspace_cpu(collector.collect_with_time(now));
        assert_eq!(keyspace_cpu.len(), REPORTED_KEYSPACE_CPU_TOP_N);
        assert!(!keyspace_cpu.contains_key(&0));
        assert!(keyspace_cpu.contains_key(&(REPORTED_KEYSPACE_CPU_TOP_N as u32)));
    }

    #[test]
    fn test_collect_reports_recent_usage_without_cumulative_spike() {
        let collector = KeyspaceCpuCollector::new();
        let first = Instant::from_timespec_second_coarse(100);
        collector.record_cpu_time(vec![(42, 5_000_000)], first);

        let keyspace_cpu = collect_keyspace_cpu(collector.collect_with_time(first));
        assert!((keyspace_cpu.get(&42).unwrap() - (5.0 / 15.0)).abs() < 1e-9);

        let later = Instant::from_timespec_second_coarse(
            100 + RECENT_KEYSPACE_CPU_USAGE_WINDOW_SECS as i64,
        );
        let keyspace_cpu = collect_keyspace_cpu(collector.collect_with_time(later));
        assert!(keyspace_cpu.is_empty());

        collector.record_cpu_time(vec![(42, 1_000_000)], later);
        let keyspace_cpu = collect_keyspace_cpu(collector.collect_with_time(later));
        assert!((keyspace_cpu.get(&42).unwrap() - (1.0 / 15.0)).abs() < 1e-9);
    }

    #[test]
    fn test_collect_sums_usage_within_sliding_window() {
        let collector = KeyspaceCpuCollector::new();
        collector.record_cpu_time(
            vec![(42, 1_000_000)],
            Instant::from_timespec_second_coarse(100),
        );
        collector.record_cpu_time(
            vec![(42, 2_000_000)],
            Instant::from_timespec_second_coarse(101),
        );

        let keyspace_cpu = collect_keyspace_cpu(
            collector.collect_with_time(Instant::from_timespec_second_coarse(101)),
        );
        assert!((keyspace_cpu.get(&42).unwrap() - 0.2).abs() < 1e-9);
    }
}