#![allow(clippy::unwrap_used)]

use lazy_static::lazy_static;
use prometheus::{
    Histogram, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec, exponential_buckets,
    linear_buckets, register_histogram, register_histogram_vec, register_int_counter_vec,
    register_int_gauge, register_int_gauge_vec,
};

// Metrics for meta server
lazy_static! {
    pub static ref SHARD_COUNT_GAUGE_VEC: IntGaugeVec =
        register_int_gauge_vec!("tici_meta_shard_count", "Number of shards", &["type"]).unwrap();
    pub static ref SHARD_BY_WRITER_COUNT_GAUGE_VEC: IntGaugeVec =
        register_int_gauge_vec!("tici_meta_shard_by_writer_count", "Number of shards on writer nodes", &["node"]).unwrap();
    pub static ref COMPACTION_BY_WRITER_COUNT_GAUGE_VEC: IntGaugeVec =
        register_int_gauge_vec!("tici_meta_compaction_by_writer_count", "Number of compactions on writer nodes", &["node"]).unwrap();
    pub static ref INDEX_COUNT_GAUGE_VEC: IntGaugeVec =
        register_int_gauge_vec!("tici_meta_index_count", "Number of indexes", &["type"]).unwrap();
    // Some events here is equivalent with SHARD_TASK_DURATION_HISTOGRAM's count, some are not though.
    pub static ref SHARD_EVENT_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tici_meta_shard_event_count",
        "Number of shard events",
        &["type"]
    )
    .unwrap();
    pub static ref WORKER_HEARTBEATS_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tici_meta_worker_heartbeats_count",
        "Number of worker heartbeats received by meta",
        &["addr"]
    )
    .unwrap();
    pub static ref SHARD_TASK_DURATION_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_shard_task_duration",
        "The time takes to handle shard tasks",
        &["type"],
        exponential_buckets(10.0, 2.0, 12).unwrap()
    )
    .unwrap();
    pub static ref SHARD_EVENT_DURATION_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_shard_event_duration",
        "The time takes to handle shard events",
        &["type"],
        exponential_buckets(0.01, 2.0, 12).unwrap()
    )
    .unwrap();
    pub static ref SHARD_SIZE_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_shard_size",
        "Bucketed histogram of approximate shard size",
        exponential_buckets(1024.0 * 1024.0, 2.0, 20).unwrap() // max bucket would be 512GB
    ).unwrap();
    pub static ref SHARD_FRAGMENTS_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_shard_fragments",
        "Bucketed histogram of fragment count per shard",
        exponential_buckets(1.0, 2.0, 10).unwrap() // 1,2,4,...,512
    )
    .unwrap();
    pub static ref BACKEND_DELAY_DURATION_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_backend_delay_duration",
        "Bucketed histogram of backend",
        &["type"],
        exponential_buckets(0.05, 2.0, 12).unwrap()
    ).unwrap();
    pub static ref SHARD_TASK_TRIGGER_REASON_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tici_meta_shard_task_trigger_reason_count",
        "Reason for triggering different shard tasks, do not filter by whether the task will succeed",
        &["reason"]
    )
    .unwrap();
    pub static ref SHARD_TASK_QUEUE_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_shard_task_queued",
        "Bucketed histogram of shard task queued.",
        &["type"],
        exponential_buckets(0.1, 2.0, 10).unwrap()  // 100ms~102.4s
    ).unwrap();
    pub static ref SHARD_TASK_BY_RUNNER_COUNTER_VEC: IntCounterVec = register_int_counter_vec!(
        "tici_meta_shard_task_by_runner_count",
        "Number of shard task per runner",
        &["id", "type"]
    )
    .unwrap();
    pub static ref SHARD_SPLIT_IMBALANCE_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_shard_split_imbalance",
        "The ratio of the max / min size of shards",
        &["type"],
        exponential_buckets(1.0, 2.0, 10).unwrap(),
    )
    .unwrap();
    pub static ref COMPACTION_RATIO_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_compaction_ratio",
        "The ratio of compaction fragment size before/after",
        &["type"],
        exponential_buckets(1.0, 2.0, 10).unwrap() // 1,2,4,...,512
    )
    .unwrap();

    pub static ref COMPACTION_INPUT_FRAGMENT_COUNT_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_compaction_input_fragment_count",
        "The number of input fragments per compaction",
        linear_buckets(1.0, 5.0, 20).unwrap() // 1,6,11,...,96
    )
    .unwrap();
    pub static ref COMPACTION_INPUT_SIZE_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_compaction_input_size",
        "Sum of input fragment sizes (bytes) per compaction, not weighted by level",
        exponential_buckets(1024.0, 4.0, 12).unwrap() // 1KiB, 4KiB, ..., 4GiB (factor 4)
    )
    .unwrap();
    pub static ref META_IS_LEADER_GAUGE: IntGauge = register_int_gauge!(
        "tici_meta_is_leader",
        "Whether this meta node believes it is leader (1) or not (0)"
    )
    .unwrap();
    pub static ref UNSCHEDULED_SHARDS_GAUGE: IntGauge = register_int_gauge!(
        "tici_unscheduled_shards_total",
        "Current number of shards without an assigned writer"
    )
    .unwrap();
}
// metrcis for shard cache
lazy_static! {
    // TODO(P1) deprecate this metric
    pub static ref SHARD_CACHE_HIT_COUNTER: IntCounterVec = register_int_counter_vec!(
        "tici_meta_shard_cache_hit_count",
        "Number of shard cache hits",
        &["type"]
    ).unwrap();
    pub static ref SHARD_CACHE_EVENT_COUNTER: IntCounterVec = register_int_counter_vec!(
        "tici_meta_shard_cache_event_count",
        "Number of shard cache event",
        &["type"]
    ).unwrap();
    pub static ref SHARD_CACHE_WARM_UP_REQ_LATENCY_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_shard_cache_warm_up_req_latency",
        "Bucketed histogram of shard cache warm up latency.",
        exponential_buckets(0.01, 2.0, 12).unwrap()  // 10ms~40.96s
    ).unwrap();
    pub static ref SHARD_CACHE_WARM_UP_QUEUE_HISTOGRAM: Histogram = register_histogram!(
        "tici_meta_shard_cache_warm_up_queued",
        "Bucketed histogram of shard cache warm up queued.",
        exponential_buckets(0.1, 2.0, 10).unwrap()  // 100ms~102.4s
    ).unwrap();
    pub static ref SHARD_CACHE_QUERY_DURATION_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "tici_meta_shard_cache_query_duration",
        "Bucketed histogram of query time.",
        &["type"],
        exponential_buckets(0.1, 2.0, 10).unwrap()  // 100ms~102.4s
    ).unwrap();
    // TODO(P1) deprecate this metric
    // NOTE: We use a gauge here to reflect the current number of shard cache allocations
    // (can increase or decrease due to split/merge/heartbeat/expire), while keeping the
    // metric name unchanged for compatibility.
    pub static ref SHARD_CACHE_ALLOCATE_COUNTER: IntGauge = register_int_gauge!(
        "tici_meta_shard_cache_allocate_count",
        "Current number of shard cache allocations"
    ).unwrap();
    pub static ref SHARD_CACHE_COUNT_GAUGE: IntGaugeVec = register_int_gauge_vec!(
        "tici_meta_shard_cache_count_gauge",
        "Current number of shard cache counts",
        &["type"]
    ).unwrap();
    // Active readers gauge
    pub static ref SHARD_CACHE_READER_COUNT_GAUGE: IntGauge = register_int_gauge!(
        "tici_meta_shard_cache_reader_count",
        "Number of active reader nodes"
    ).unwrap();
    // TODO(P1) deprecate this metric
    // Reader module shard count (event-driven)
    pub static ref READER_SHARD_COUNT_GAUGE: IntGauge = register_int_gauge!(
        "tici_meta_reader_shard_count",
        "Total shards tracked by reader module (event-driven)."
    ).unwrap();
}
// metrcis for import into
lazy_static! {
    pub static ref IMPORT_PARTITION_UPLOAD_COUNTER: IntCounterVec = register_int_counter_vec!(
        "tici_meta_import_partition_upload_finish_count",
        "Number of import partition upload finish requests",
        &["result"]
    )
    .unwrap();
    pub static ref IMPORT_INDEX_UPLOAD_COUNTER: IntCounterVec = register_int_counter_vec!(
        "tici_meta_import_index_upload_finish_count",
        "Number of import index upload finish requests",
        &["result"]
    )
    .unwrap();
    pub static ref IMPORT_PARTITION_BUILD_COUNTER: IntCounterVec = register_int_counter_vec!(
        "tici_meta_import_partition_build_finish_count",
        "Number of import partition build finish requests",
        &["result"]
    )
    .unwrap();
}
/// rust-prometheus doesn't support clear existing data.
pub fn clear_prom_metrics() {
    INDEX_COUNT_GAUGE_VEC.reset();
    SHARD_COUNT_GAUGE_VEC.reset();
    SHARD_EVENT_COUNTER_VEC.reset();
    WORKER_HEARTBEATS_COUNTER_VEC.reset();
    META_IS_LEADER_GAUGE.set(0);
    UNSCHEDULED_SHARDS_GAUGE.set(0);

    SHARD_CACHE_EVENT_COUNTER.reset();
    SHARD_CACHE_QUERY_DURATION_HISTOGRAM.reset();
    SHARD_CACHE_COUNT_GAUGE.reset();

    IMPORT_PARTITION_UPLOAD_COUNTER.reset();
    IMPORT_INDEX_UPLOAD_COUNTER.reset();
    IMPORT_PARTITION_BUILD_COUNTER.reset();
}

pub fn update_meta_leader_metric(is_leader: bool) {
    let value = if is_leader { 1 } else { 0 };
    META_IS_LEADER_GAUGE.set(value);
}

#[cfg(test)]
mod tests {
    use super::{META_IS_LEADER_GAUGE, UNSCHEDULED_SHARDS_GAUGE, update_meta_leader_metric};

    #[test]
    fn meta_leader_metric_tracks_state() {
        update_meta_leader_metric(true);
        assert_eq!(META_IS_LEADER_GAUGE.get(), 1);

        update_meta_leader_metric(false);
        assert_eq!(META_IS_LEADER_GAUGE.get(), 0);
    }

    #[test]
    fn unscheduled_shards_metric_is_a_gauge() {
        let before = UNSCHEDULED_SHARDS_GAUGE.get();

        UNSCHEDULED_SHARDS_GAUGE.inc();
        assert_eq!(UNSCHEDULED_SHARDS_GAUGE.get(), before + 1);

        UNSCHEDULED_SHARDS_GAUGE.dec();
        assert_eq!(UNSCHEDULED_SHARDS_GAUGE.get(), before);
    }
}