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

use std::sync::OnceLock;

use lazy_static::lazy_static;
use prometheus::*;
use prometheus_static_metric::*;

make_static_metric! {
    pub label_enum PDReconnectEventKind {
        success,
        failure,
        no_need,
        cancel,
        try_connect,
    }

    pub struct PDReconnectEventCounterVec: IntCounter {
        "type" => PDReconnectEventKind,
    }
}

lazy_static! {
    pub static ref PD_REQUEST_HISTOGRAM_VEC: HistogramVec = register_histogram_vec!(
        "tikv_pd_request_duration_seconds",
        "Bucketed histogram of PD requests duration",
        &["type"]
    )
    .unwrap();
    pub static ref PD_HEARTBEAT_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tikv_pd_heartbeat_message_total",
        "Total number of PD heartbeat messages.",
        &["type"]
    )
    .unwrap();
    pub static ref PD_BUCKETS_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tikv_pd_buckets_message_total",
        "Total number of PD buckets messages.",
        &["type"]
    )
    .unwrap();
    pub static ref PD_RECONNECT_COUNTER_VEC: PDReconnectEventCounterVec =
        register_static_int_counter_vec!(
            PDReconnectEventCounterVec,
            "tikv_pd_reconnect_total",
            "Total number of PD reconnections.",
            &["type"]
        )
        .unwrap();
    pub static ref PD_PENDING_HEARTBEAT_GAUGE: IntGauge = register_int_gauge!(
        "tikv_pd_pending_heartbeat_total",
        "Total number of pending region heartbeat"
    )
    .unwrap();
    pub static ref PD_PENDING_BUCKETS_GAUGE: IntGauge = register_int_gauge!(
        "tikv_pd_pending_buckets_total",
        "Total number of pending region buckets"
    )
    .unwrap();
    pub static ref PD_VALIDATE_PEER_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tikv_pd_validate_peer_total",
        "Total number of pd worker validate peer task.",
        &["type"]
    )
    .unwrap();
    pub static ref REGION_READ_KEYS_HISTOGRAM: Histogram = register_histogram!(
        "tikv_region_read_keys",
        "Histogram of keys written for regions",
        exponential_buckets(1.0, 2.0, 20).unwrap()
    )
    .unwrap();
    pub static ref STORE_ENGINE_SIZE_GAUGE_VEC: IntGaugeVec = register_int_gauge_vec!(
        "tikv_engine_size_bytes",
        "Sizes of each column families",
        &["db", "type"]
    )
    .unwrap();
    pub static ref STORE_ENGINE_MEM_SIZE_GAUGE_VEC: IntGaugeVec = register_int_gauge_vec!(
        "kv_engine_mem_size_bytes",
        "Sizes of memory used by engine",
        &["db", "type"]
    )
    .unwrap();
    pub static ref STORE_ENGINE_MEM_QUOTA_GAUGE: Gauge = register_gauge!(
        "kv_engine_mem_size_quota_bytes",
        "Quota of memory usage of memtables in KvEngine"
    )
    .unwrap();
    pub static ref STORE_ENGINE_FLOW_VEC: IntCounterVec = register_int_counter_vec!(
        "tikv_engine_flow_bytes",
        "Bytes and keys of read/written",
        &["db", "type"]
    )
    .unwrap();
    pub static ref REGION_READ_BYTES_HISTOGRAM: Histogram = register_histogram!(
        "tikv_region_read_bytes",
        "Histogram of bytes written for regions",
        exponential_buckets(256.0, 2.0, 20).unwrap()
    )
    .unwrap();
    pub static ref REGION_WRITTEN_BYTES_HISTOGRAM: Histogram = register_histogram!(
        "tikv_region_written_bytes",
        "Histogram of bytes written for regions",
        exponential_buckets(256.0, 2.0, 20).unwrap()
    )
    .unwrap();
    pub static ref REGION_WRITTEN_KEYS_HISTOGRAM: Histogram = register_histogram!(
        "tikv_region_written_keys",
        "Histogram of keys written for regions",
        exponential_buckets(1.0, 2.0, 20).unwrap()
    )
    .unwrap();
    pub static ref REQUEST_FORWARDED_GAUGE_VEC: IntGaugeVec = register_int_gauge_vec!(
        "tikv_pd_request_forwarded",
        "The status to indicate if the request is forwarded",
        &["host"]
    )
    .unwrap();
    pub static ref PD_PENDING_TSO_REQUEST_GAUGE: IntGauge = register_int_gauge!(
        "tikv_pd_pending_tso_request_total",
        "Total number of pending tso requests"
    )
    .unwrap();
}

// Storage-size metrics expose keyspace labels differently by product:
// Starter/Essential keep `keyspace_id`; Premium expects `keyspace_name`.
// Prometheus forbids re-registering the same metric name with a different
// descriptor, so each schema is chosen once at startup and frozen. Mode and
// collector live behind a single OnceLock per metric family to avoid any
// chance of the emitter mode disagreeing with what Prometheus registered.
struct StoreSizeMetric {
    use_keyspace_name: bool,
    vec: IntGaugeVec,
}

static STORE_SIZE_METRIC: OnceLock<StoreSizeMetric> = OnceLock::new();

struct KeyspaceSizeMetric {
    use_keyspace_name: bool,
    vec: IntGaugeVec,
}

static KEYSPACE_SIZE_METRIC: OnceLock<KeyspaceSizeMetric> = OnceLock::new();

// Starter/Essential: keyspace_id label schema (legacy).
const STORE_SIZE_LABELS_KEYSPACE_ID: &[&str] =
    &["type", "region_id", "keyspace_id", "storage_class"];
// Premium: keyspace_name label schema.
const STORE_SIZE_LABELS_KEYSPACE_NAME: &[&str] =
    &["type", "region_id", "keyspace_name", "storage_class"];

const KEYSPACE_SIZE_LABELS_KEYSPACE_ID: &[&str] = &["type", "keyspace_id", "storage_class"];
const KEYSPACE_SIZE_LABELS_KEYSPACE_NAME: &[&str] = &["type", "keyspace_name", "storage_class"];

fn register_store_size(use_keyspace_name: bool) -> StoreSizeMetric {
    let labels = if use_keyspace_name {
        STORE_SIZE_LABELS_KEYSPACE_NAME
    } else {
        STORE_SIZE_LABELS_KEYSPACE_ID
    };
    StoreSizeMetric {
        use_keyspace_name,
        vec: register_int_gauge_vec!("tikv_store_size_bytes", "Size of storage.", labels).unwrap(),
    }
}

fn register_keyspace_size(use_keyspace_name: bool) -> KeyspaceSizeMetric {
    let labels = if use_keyspace_name {
        KEYSPACE_SIZE_LABELS_KEYSPACE_NAME
    } else {
        KEYSPACE_SIZE_LABELS_KEYSPACE_ID
    };
    KeyspaceSizeMetric {
        use_keyspace_name,
        vec: register_int_gauge_vec!(
            "tikv_keyspace_size_bytes",
            "Keyspace-level size of storage.",
            labels
        )
        .unwrap(),
    }
}

/// Initialize `tikv_store_size_bytes` with the configured label schema. Must
/// be called before any emit. Calling twice with the same mode is a no-op;
/// calling with a different mode after init panics, because Prometheus has
/// already frozen the descriptor.
pub fn init_store_size_metric(use_keyspace_name: bool) {
    let actual = STORE_SIZE_METRIC.get_or_init(|| register_store_size(use_keyspace_name));
    assert_eq!(
        actual.use_keyspace_name, use_keyspace_name,
        "tikv_store_size_bytes already initialized with a different schema"
    );
}

/// Initialize `tikv_keyspace_size_bytes` with the configured label schema.
/// Must be called before any emit. Calling twice with the same mode is a
/// no-op; calling with a different mode after init panics, because Prometheus
/// has already frozen the descriptor.
pub fn init_keyspace_size_metric(use_keyspace_name: bool) {
    let actual = KEYSPACE_SIZE_METRIC.get_or_init(|| register_keyspace_size(use_keyspace_name));
    assert_eq!(
        actual.use_keyspace_name, use_keyspace_name,
        "tikv_keyspace_size_bytes already initialized with a different schema"
    );
}

/// Return whether `tikv_store_size_bytes` uses `keyspace_name`.
pub fn store_size_uses_keyspace_name() -> bool {
    STORE_SIZE_METRIC.get().is_some_and(|m| m.use_keyspace_name)
}

/// Return whether `tikv_keyspace_size_bytes` uses `keyspace_name`.
pub fn keyspace_size_uses_keyspace_name() -> bool {
    KEYSPACE_SIZE_METRIC
        .get()
        .is_some_and(|m| m.use_keyspace_name)
}

/// Access the `tikv_store_size_bytes` collector. Falls back to the legacy
/// `keyspace_id` schema if [`init_store_size_metric`] was never called (test
/// paths that don't go through cloud_server startup).
pub fn store_size_gauge_vec() -> &'static IntGaugeVec {
    &STORE_SIZE_METRIC
        .get_or_init(|| register_store_size(false))
        .vec
}

/// Access the `tikv_keyspace_size_bytes` collector. Falls back to the legacy
/// `keyspace_id` schema if [`init_keyspace_size_metric`] was never called
/// (test paths that don't go through cloud_server startup).
pub fn keyspace_size_gauge_vec() -> &'static IntGaugeVec {
    &KEYSPACE_SIZE_METRIC
        .get_or_init(|| register_keyspace_size(false))
        .vec
}

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

    // Prometheus's global registry plus `STORE_SIZE_METRIC` is process-wide, so
    // all assertions about init behavior must run inside one #[test] to keep
    // ordering deterministic and avoid cross-test contamination.
    #[test]
    fn storage_size_metric_init_is_locked_after_first_call() {
        init_store_size_metric(false);
        init_keyspace_size_metric(false);
        assert!(!store_size_uses_keyspace_name());
        assert!(!keyspace_size_uses_keyspace_name());
        store_size_gauge_vec()
            .get_metric_with_label_values(&["used", "1", "5", "standard"])
            .unwrap();
        keyspace_size_gauge_vec()
            .get_metric_with_label_values(&["used", "5", "standard"])
            .unwrap();

        // Same mode is idempotent.
        init_store_size_metric(false);
        init_keyspace_size_metric(false);

        // Switching modes after registration must fail loudly; Prometheus has
        // already frozen the descriptor and emitters would otherwise drift.
        let res = std::panic::catch_unwind(|| init_store_size_metric(true));
        assert!(res.is_err(), "switching schema after init must panic");
        let res = std::panic::catch_unwind(|| init_keyspace_size_metric(true));
        assert!(res.is_err(), "switching schema after init must panic");
        assert!(!store_size_uses_keyspace_name());
        assert!(!keyspace_size_uses_keyspace_name());
    }
}