use lazy_static::lazy_static;
use prometheus::{HistogramVec, IntCounterVec, IntGaugeVec, exponential_buckets, linear_buckets};
use tici_common::metrics::{
    OperationMetrics, OperationMetricsGuard, begin_operation_metrics0, begin_operation_metrics1,
    register_histogram_with_labels, register_histogram_with_labels_and_buckets,
    register_int_counter_with_labels, register_int_gauge_with_labels,
};
use tici_grpc::{CompactFragRequest, FinishCompactFragRequest, FragMeta};

lazy_static! {
    // Compaction metrics
    pub static ref WORKER_COMPACTION_DURATION_HISTOGRAM: HistogramVec =
        register_histogram_with_labels_and_buckets(
            "tici_worker_compaction_duration",
            "The duration of worker compaction operations",
            &["stage"],
            exponential_buckets(0.1, 2.0, 16).unwrap(),
        );
    pub static ref WORKER_COMPACTION_COUNTER: IntCounterVec = register_int_counter_with_labels(
        "tici_worker_compaction_count",
        "The count of worker compaction operations",
        &["stage", "result"],
    );
    pub static ref WORKER_COMPACTION_INFLIGHT_GAUGE: IntGaugeVec = register_int_gauge_with_labels(
        "tici_worker_compaction_inflight",
        "The inflight count of worker compaction operations",
        &["stage"],
    );
    pub static ref WORKER_COMPACTION_FRAGMENT_COUNT_HISTOGRAM: HistogramVec =
        register_histogram_with_labels_and_buckets(
            "tici_worker_compaction_fragment_count",
            "The number of fragments per worker compaction",
            &["name"],
            linear_buckets(1.0, 5.0, 100).unwrap(),
        );
    pub static ref WORKER_COMPACTION_SIZE_HISTOGRAM: HistogramVec =
        register_histogram_with_labels_and_buckets(
            "tici_worker_compaction_fragment_size",
            "The total fragment size in bytes per worker compaction",
            &["name"],
            exponential_buckets(1024.0, 4.0, 12).unwrap(),
        );
    // Dedup metrics
    pub static ref DEDUP_RUN_COUNTER: IntCounterVec = register_int_counter_with_labels(
        "tici_dedup_run_count",
        "The count of dedup operations",
        &["result"],
    );
    pub static ref DEDUP_IN_PROGRESS_GAUGE: IntGaugeVec = register_int_gauge_with_labels(
        "tici_dedup_in_progress",
        "The inflight count of dedup operations",
        &[],
    );
    pub static ref DEDUP_STAGE_LATENCY: HistogramVec = register_histogram_with_labels(
        "tici_dedup_stage_latency",
        "Latency of compactor dedup stages (s)",
        &["stage"],
    );
}

pub type DedupMetricsGuard = OperationMetricsGuard;

pub fn begin_dedup_metrics() -> DedupMetricsGuard {
    begin_operation_metrics0(OperationMetrics::new_without_duration(
        &DEDUP_RUN_COUNTER,
        &DEDUP_IN_PROGRESS_GAUGE,
    ))
}

pub fn begin_worker_compaction_metrics(phase: &'static str) -> OperationMetricsGuard {
    begin_operation_metrics1(
        OperationMetrics::new(
            &WORKER_COMPACTION_DURATION_HISTOGRAM,
            &WORKER_COMPACTION_COUNTER,
            &WORKER_COMPACTION_INFLIGHT_GAUGE,
        ),
        phase,
    )
}

pub trait CompactionMetricsRequest {
    fn frags(&self) -> &[FragMeta];
}

impl CompactionMetricsRequest for CompactFragRequest {
    fn frags(&self) -> &[FragMeta] {
        &self.input_frags
    }
}

impl CompactionMetricsRequest for FinishCompactFragRequest {
    fn frags(&self) -> &[FragMeta] {
        &self.output_frags
    }
}

pub fn observe_compaction<T: CompactionMetricsRequest>(req: &T, name: &str) {
    let frags = req.frags();
    WORKER_COMPACTION_FRAGMENT_COUNT_HISTOGRAM
        .with_label_values(&[name])
        .observe(frags.len() as f64);
    WORKER_COMPACTION_SIZE_HISTOGRAM
        .with_label_values(&[name])
        .observe(sum_frag_sizes(frags) as f64);
}

fn sum_frag_sizes(frags: &[FragMeta]) -> u64 {
    frags
        .iter()
        .filter_map(|frag| frag.property.as_ref().map(|property| property.size))
        .sum()
}