// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{collections::HashMap, io::Write};

use kvproto::{errorpb, pdpb};
use lazy_static::lazy_static;
use prometheus::{proto::MetricType, *};
use prometheus_static_metric::*;

#[cfg(target_os = "linux")]
mod threads_linux;
#[cfg(target_os = "linux")]
pub use self::threads_linux::{ThreadInfoStatistics, monitor_threads};

#[cfg(target_os = "linux")]
mod process_linux;
#[cfg(target_os = "linux")]
pub use self::process_linux::monitor_process;

#[cfg(target_os = "linux")]
mod psi_linux;
#[cfg(target_os = "linux")]
pub use self::psi_linux::{PSI_PATH, monitor_system_psi};

#[cfg(not(target_os = "linux"))]
mod threads_dummy;
#[cfg(not(target_os = "linux"))]
pub use self::threads_dummy::{ThreadInfoStatistics, monitor_threads};

#[cfg(not(target_os = "linux"))]
mod process_dummy;
pub use self::allocator_metrics::monitor_allocator_stats;
#[cfg(not(target_os = "linux"))]
pub use self::process_dummy::monitor_process;

#[cfg(not(target_os = "linux"))]
mod psi_dummy;
#[cfg(not(target_os = "linux"))]
pub use self::psi_dummy::monitor_system_psi;
#[cfg(not(target_os = "linux"))]
pub const PSI_PATH: &str = "";

pub mod allocator_metrics;

pub use self::metrics_reader::HistogramReader;

mod metrics_reader;

pub type RecordPairVec = Vec<pdpb::RecordPair>;

/// Bounded categories for `ServerIsBusy` responses constructed for clients.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ServerIsBusyType {
    Backup,
    CdcIncrementalScan,
    CoprocessorCpuThrottled,
    CoprocessorMaxPendingTasks,
    CoprocessorMemoryQuota,
    Deadline,
    GcWorker,
    Raftstore,
    RaftstoreIngest,
    ReplicationWorker,
    Scheduler,
    SstImporter,
    SstSnapshot,
    TxnFile,
}

impl ServerIsBusyType {
    fn as_str(self) -> &'static str {
        match self {
            Self::Backup => "backup",
            Self::CdcIncrementalScan => "cdc_incremental_scan",
            Self::CoprocessorCpuThrottled => "coprocessor_cpu_throttled",
            Self::CoprocessorMaxPendingTasks => "coprocessor_max_pending_tasks",
            Self::CoprocessorMemoryQuota => "coprocessor_memory_quota",
            Self::Deadline => "deadline",
            Self::GcWorker => "gc_worker",
            Self::Raftstore => "raftstore",
            Self::RaftstoreIngest => "raftstore_ingest",
            Self::ReplicationWorker => "replication_worker",
            Self::Scheduler => "scheduler",
            Self::SstImporter => "sst_importer",
            Self::SstSnapshot => "sst_snapshot",
            Self::TxnFile => "txn_file",
        }
    }
}

pub fn dump(should_simplify: bool) -> String {
    let mut buffer = vec![];
    dump_to(&mut buffer, should_simplify);
    String::from_utf8(buffer).unwrap()
}

pub fn dump_to(w: &mut impl Write, should_simplify: bool) {
    let encoder = TextEncoder::new();
    let metric_families = prometheus::gather();
    if !should_simplify {
        if let Err(e) = encoder.encode(&metric_families, w) {
            warn!("prometheus encoding error"; "err" => ?e);
        }
        return;
    }

    // filter out mertics that has no sample values
    for mut mf in metric_families {
        let mut metrics = mf.take_metric().into_vec();
        match mf.get_field_type() {
            MetricType::COUNTER => {
                metrics.retain(|m| m.get_counter().get_value() > 0.0);
            }
            MetricType::HISTOGRAM => metrics.retain(|m| m.get_histogram().get_sample_count() > 0),
            _ => {}
        }
        if !metrics.is_empty() {
            mf.set_metric(metrics.into());
            if let Err(e) = encoder.encode(&[mf], w) {
                warn!("prometheus encoding error"; "err" => ?e);
            }
        }
    }
}

make_auto_flush_static_metric! {
    // Some non-txn related types are placed here.
    // ref `TXN_COMMAND_THROTTLE_TIME_COUNTER_VEC`.
    pub label_enum ThrottleType {
        dag,
        analyze_full_sampling,
        quota_limiter_auto_tuned,
    }

    pub struct NonTxnCommandThrottleTimeCounterVec: LocalIntCounter {
        "type" => ThrottleType,
    }
}

lazy_static! {
    pub static ref CRITICAL_ERROR: IntCounterVec = register_int_counter_vec!(
        "tikv_critical_error_total",
        "Counter of critical error.",
        &["type"]
    )
    .unwrap();
    pub static ref FAIL_CLOSE_COUNTER: IntCounterVec =
        register_int_counter_vec!(
            "tikv_fail_close_total",
            "Total number of fail-close events where TiKV handled an unexpected but non-data-corrupting state by returning an error instead of panicking.",
            &["component", "operation", "category", "reason", "action"]
        )
        .unwrap();
    static ref SERVER_IS_BUSY_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tikv_server_is_busy_total",
        "Total number of ServerIsBusy responses constructed for clients.",
        &["type"]
    )
    .unwrap();
    pub static ref NON_TXN_COMMAND_THROTTLE_TIME_COUNTER_VEC: IntCounterVec =
        register_int_counter_vec!(
            "tikv_non_txn_command_throttle_time_total",
            "Total throttle time (microsecond) of non txn processing.",
            &["type"]
        )
        .unwrap();
    pub static ref NON_TXN_COMMAND_THROTTLE_TIME_COUNTER_VEC_STATIC:
        NonTxnCommandThrottleTimeCounterVec = auto_flush_from!(
            NON_TXN_COMMAND_THROTTLE_TIME_COUNTER_VEC,
            NonTxnCommandThrottleTimeCounterVec
        );
    pub static ref INSTANCE_BACKEND_CPU_QUOTA: IntGauge =
        register_int_gauge!("tikv_backend_cpu_quota", "cpu quota for backend request").unwrap();
}

/// Constructs a `ServerIsBusy` error and records it as a response intended for
/// a client.
pub fn new_server_is_busy(
    type_: ServerIsBusyType,
    reason: impl Into<String>,
) -> errorpb::ServerIsBusy {
    SERVER_IS_BUSY_COUNTER_VEC
        .with_label_values(&[type_.as_str()])
        .inc();
    let mut error = errorpb::ServerIsBusy::default();
    error.set_reason(reason.into());
    error
}

pub fn convert_record_pairs(m: HashMap<String, u64>) -> RecordPairVec {
    m.into_iter()
        .map(|(k, v)| {
            let mut pair = pdpb::RecordPair::default();
            pair.set_key(k);
            pair.set_value(v);
            pair
        })
        .collect()
}

#[cfg(target_os = "linux")]
pub fn warn_if_kernel_metrics_disabled() {
    use std::{fs, path::Path};

    fn read_flag(path: &str) -> Option<String> {
        fs::read_to_string(path).ok().map(|s| s.trim().to_owned())
    }

    match read_flag("/proc/sys/kernel/task_delayacct") {
        Some(v) if v == "1" => {}
        Some(_) => {
            warn!(
                "task_delayacct is disabled; thread_blkio_delay_seconds_total will stay zero";
                "hint" => "echo 1 > /proc/sys/kernel/task_delayacct"
            );
        }
        None => {
            warn!(
                "task_delayacct sysctl missing; thread_blkio_delay_seconds_total will stay zero";
                "path" => "/proc/sys/kernel/task_delayacct"
            );
        }
    }

    match read_flag("/proc/sys/kernel/sched_schedstats") {
        Some(v) if v == "1" => {}
        Some(_) => {
            warn!(
                "sched_schedstats is disabled; run_queue_wait and sched_timeslices metrics will stay zero";
                "hint" => "echo 1 > /proc/sys/kernel/sched_schedstats or boot with schedstats=enable"
            );
        }
        None => {
            warn!(
                "sched_schedstats sysctl missing; run_queue_wait and sched_timeslices metrics will stay zero";
                "path" => "/proc/sys/kernel/sched_schedstats"
            );
        }
    }

    if !Path::new("/proc/pressure/cpu").exists() {
        warn!(
            "PSI /proc/pressure not available; pressure_stall metrics will be empty";
            "hint" => "requires CONFIG_PSI=y and a kernel that exposes /proc/pressure"
        );
    }

    // cgroup v2 preferred for accurate PSI in containerized environments.
    if !Path::new("/sys/fs/cgroup/cgroup.controllers").exists() {
        warn!(
            "cgroup v2 not detected; PSI accuracy may be reduced in containers";
            "hint" => "enable cgroup v2 (unified hierarchy) on the host"
        );
    }
}

#[cfg(not(target_os = "linux"))]
pub fn warn_if_kernel_metrics_disabled() {}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use prometheus::*;

    use super::*;

    #[test]
    fn test_dump_metrics() {
        // register some metrics
        let _counter = register_int_counter!("test_counter", "this is a counter for test").unwrap();
        let counter_vec =
            register_int_counter_vec!("test_counter_vec", "test counter vec", &["label"]).unwrap();
        let histogram = register_histogram_vec!(
            "test_histogram",
            "test histogram",
            &["type"],
            exponential_buckets(0.01, 2.0, 20).unwrap()
        )
        .unwrap();
        let gauge = register_gauge!("test_gauge", "test gauge").unwrap();

        fn check_duplicate(s: &str) {
            let mut lines = HashSet::new();
            for l in s.lines() {
                assert!(lines.insert(l));
            }
        }

        // test all data is 0.
        let full_metrics = dump(false);
        assert!(!full_metrics.is_empty());
        check_duplicate(&full_metrics);

        let filtered_metrics = dump(true);
        check_duplicate(&full_metrics);
        assert!(full_metrics.len() > filtered_metrics.len());

        counter_vec.with_label_values(&["test"]).inc();
        histogram.with_label_values(&["test"]).observe(1.0);
        gauge.inc();

        let new_full_metrics = dump(false);
        assert!(!new_full_metrics.is_empty());
        check_duplicate(&new_full_metrics);
        assert!(new_full_metrics.len() > full_metrics.len());

        let new_filtered_metrics = dump(true);
        check_duplicate(&new_filtered_metrics);
        assert!(new_filtered_metrics.len() > filtered_metrics.len());
        assert!(new_full_metrics.len() > new_filtered_metrics.len());
    }

    #[test]
    fn test_new_server_is_busy() {
        let counter = SERVER_IS_BUSY_COUNTER_VEC.with_label_values(&["scheduler"]);
        let before = counter.get();

        let error = new_server_is_busy(ServerIsBusyType::Scheduler, "scheduler is busy");

        assert_eq!(counter.get(), before + 1);
        assert_eq!(error.reason, "scheduler is busy");
    }
}