use std::{
    fmt::Display,
    sync::atomic::{AtomicBool, Ordering},
    time::Instant,
};

use prometheus::{
    Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, register_histogram,
    register_histogram_vec, register_int_counter, register_int_counter_vec, register_int_gauge,
    register_int_gauge_vec,
};

fn handle_register_result<T, E: Display>(name: &str, result: Result<T, E>) -> T {
    result.unwrap_or_else(|err| panic!("failed to register prometheus metric {name}: {err}"))
}

pub fn register_histogram_with_labels_and_buckets(
    name: &str,
    help: &str,
    labels: &[&str],
    buckets: Vec<f64>,
) -> HistogramVec {
    let res = register_histogram_vec!(name, help, labels, buckets);
    handle_register_result(name, res)
}

pub fn register_histogram_with_labels(name: &str, help: &str, labels: &[&str]) -> HistogramVec {
    let res = register_histogram_vec!(name, help, labels);
    handle_register_result(name, res)
}

pub fn register_histogram_with_buckets(name: &str, help: &str, buckets: Vec<f64>) -> Histogram {
    let res = register_histogram!(name, help, buckets);
    handle_register_result(name, res)
}

pub fn register_int_counter_with_labels(name: &str, help: &str, labels: &[&str]) -> IntCounterVec {
    let res = register_int_counter_vec!(name, help, labels);
    handle_register_result(name, res)
}

pub fn register_int_counter(name: &str, help: &str) -> IntCounter {
    let res = register_int_counter!(name, help);
    handle_register_result(name, res)
}

pub fn register_int_gauge_with_labels(name: &str, help: &str, labels: &[&str]) -> IntGaugeVec {
    let res = register_int_gauge_vec!(name, help, labels);
    handle_register_result(name, res)
}

pub fn register_int_gauge(name: &str, help: &str) -> IntGauge {
    let res = register_int_gauge!(name, help);
    handle_register_result(name, res)
}

const RESULT_OK: &str = "ok";
const RESULT_ERROR: &str = "error";
const MAX_OPERATION_LABEL_COUNT: usize = 3;
const MAX_OPERATION_COUNTER_LABEL_COUNT: usize = MAX_OPERATION_LABEL_COUNT + 1;

#[derive(Clone, Copy)]
pub struct OperationMetrics {
    duration_histogram: Option<&'static HistogramVec>,
    counter: &'static IntCounterVec,
    inflight_gauge: &'static IntGaugeVec,
}

impl OperationMetrics {
    pub const fn new(
        duration_histogram: &'static HistogramVec,
        counter: &'static IntCounterVec,
        inflight_gauge: &'static IntGaugeVec,
    ) -> Self {
        Self {
            duration_histogram: Some(duration_histogram),
            counter,
            inflight_gauge,
        }
    }

    pub const fn new_without_duration(
        counter: &'static IntCounterVec,
        inflight_gauge: &'static IntGaugeVec,
    ) -> Self {
        Self {
            duration_histogram: None,
            counter,
            inflight_gauge,
        }
    }
}

pub struct OperationMetricsGuard {
    metrics: OperationMetrics,
    labels: OperationMetricLabels,
    watch: Instant,
    has_error: AtomicBool,
}

impl OperationMetricsGuard {
    fn new(metrics: OperationMetrics, labels: OperationMetricLabels) -> Self {
        metrics
            .inflight_gauge
            .with_label_values(labels.as_slice())
            .inc();
        Self {
            metrics,
            labels,
            watch: Instant::now(),
            has_error: AtomicBool::new(false),
        }
    }

    pub fn mark_error(&self) {
        self.has_error.store(true, Ordering::Relaxed);
    }

    pub fn mark_result<T, E>(&self, result: &std::result::Result<T, E>) {
        if result.is_err() {
            self.mark_error();
        }
    }
}

impl Drop for OperationMetricsGuard {
    fn drop(&mut self) {
        self.metrics
            .inflight_gauge
            .with_label_values(self.labels.as_slice())
            .dec();
        if let Some(duration_histogram) = self.metrics.duration_histogram {
            duration_histogram
                .with_label_values(self.labels.as_slice())
                .observe(self.watch.elapsed().as_secs_f64());
        }

        let result = if self.has_error.load(Ordering::Relaxed) {
            RESULT_ERROR
        } else {
            RESULT_OK
        };
        let counter_label_values = self.labels.counter_label_values(result);
        self.metrics
            .counter
            .with_label_values(&counter_label_values[..self.labels.counter_label_count()])
            .inc();
    }
}

#[derive(Clone, Copy)]
struct OperationMetricLabels {
    values: [&'static str; MAX_OPERATION_LABEL_COUNT],
    len: usize,
}

impl OperationMetricLabels {
    const fn new0() -> Self {
        Self {
            values: ["", "", ""],
            len: 0,
        }
    }

    const fn new1(label0: &'static str) -> Self {
        Self {
            values: [label0, "", ""],
            len: 1,
        }
    }

    const fn new2(label0: &'static str, label1: &'static str) -> Self {
        Self {
            values: [label0, label1, ""],
            len: 2,
        }
    }

    const fn new3(label0: &'static str, label1: &'static str, label2: &'static str) -> Self {
        Self {
            values: [label0, label1, label2],
            len: 3,
        }
    }

    fn as_slice(&self) -> &[&'static str] {
        &self.values[..self.len]
    }

    fn counter_label_values(
        &self,
        result: &'static str,
    ) -> [&'static str; MAX_OPERATION_COUNTER_LABEL_COUNT] {
        let mut values = [""; MAX_OPERATION_COUNTER_LABEL_COUNT];
        values[..self.len].copy_from_slice(self.as_slice());
        values[self.len] = result;
        values
    }

    const fn counter_label_count(&self) -> usize {
        self.len + 1
    }
}

pub fn begin_operation_metrics(
    metrics: OperationMetrics,
    scope: &'static str,
    metric_type: &'static str,
) -> OperationMetricsGuard {
    OperationMetricsGuard::new(metrics, OperationMetricLabels::new2(scope, metric_type))
}

pub fn begin_operation_metrics0(metrics: OperationMetrics) -> OperationMetricsGuard {
    OperationMetricsGuard::new(metrics, OperationMetricLabels::new0())
}

pub fn begin_operation_metrics1(
    metrics: OperationMetrics,
    label0: &'static str,
) -> OperationMetricsGuard {
    OperationMetricsGuard::new(metrics, OperationMetricLabels::new1(label0))
}

pub fn begin_operation_metrics3(
    metrics: OperationMetrics,
    label0: &'static str,
    label1: &'static str,
    label2: &'static str,
) -> OperationMetricsGuard {
    OperationMetricsGuard::new(metrics, OperationMetricLabels::new3(label0, label1, label2))
}

#[cfg(test)]
mod tests {
    use prometheus::{HistogramOpts, HistogramVec, IntCounterVec, IntGaugeVec, Opts};

    use super::{OperationMetrics, begin_operation_metrics, begin_operation_metrics3};

    fn make_test_metrics(label_names: &[&str]) -> OperationMetrics {
        let histogram = Box::leak(Box::new(
            HistogramVec::new(
                HistogramOpts::new("test_duration", "test duration"),
                label_names,
            )
            .unwrap(),
        ));
        let counter_label_names: Vec<&str> = label_names
            .iter()
            .copied()
            .chain(std::iter::once("result"))
            .collect();
        let counter = Box::leak(Box::new(
            IntCounterVec::new(Opts::new("test_count", "test count"), &counter_label_names)
                .unwrap(),
        ));
        let inflight = Box::leak(Box::new(
            IntGaugeVec::new(Opts::new("test_inflight", "test inflight"), label_names).unwrap(),
        ));
        OperationMetrics::new(histogram, counter, inflight)
    }

    #[test]
    fn test_operation_metrics_guard_ok_with_two_labels() {
        let metrics = make_test_metrics(&["scope", "type"]);
        let inflight = metrics
            .inflight_gauge
            .with_label_values(&["ut_scope", "ut_type"]);
        let counter = metrics
            .counter
            .with_label_values(&["ut_scope", "ut_type", "ok"]);
        let before_inflight = inflight.get();
        let before_counter = counter.get();

        {
            let _guard = begin_operation_metrics(metrics, "ut_scope", "ut_type");
            assert_eq!(inflight.get(), before_inflight + 1);
        }

        assert_eq!(inflight.get(), before_inflight);
        assert_eq!(counter.get(), before_counter + 1);
    }

    #[test]
    fn test_operation_metrics_guard_error_with_three_labels() {
        let metrics = make_test_metrics(&["service", "backend", "type"]);
        let inflight =
            metrics
                .inflight_gauge
                .with_label_values(&["ut_service", "ut_backend", "ut_type"]);
        let counter =
            metrics
                .counter
                .with_label_values(&["ut_service", "ut_backend", "ut_type", "error"]);
        let before_inflight = inflight.get();
        let before_counter = counter.get();

        {
            let guard = begin_operation_metrics3(metrics, "ut_service", "ut_backend", "ut_type");
            assert_eq!(inflight.get(), before_inflight + 1);
            guard.mark_error();
        }

        assert_eq!(inflight.get(), before_inflight);
        assert_eq!(counter.get(), before_counter + 1);
    }
}