//! Prometheus metrics registry for partition nodes.

#[cfg(feature = "distributed")]
use prometheus_client::{
    encoding::text::encode,
    metrics::{counter::Counter, family::Family, histogram::Histogram},
    registry::Registry,
};
#[cfg(feature = "distributed")]
use std::sync::{Arc, Mutex};

/// Label set for query metrics.
#[cfg(feature = "distributed")]
#[derive(Clone, Hash, PartialEq, Eq, prometheus_client::encoding::EncodeLabelSet, Debug)]
pub struct QueryLabels {
    pub backend: String,
    pub status: String, // "ok" | "error"
}

/// Prometheus metrics for a partition node.
#[cfg(feature = "distributed")]
pub struct PartitionNodeMetrics {
    pub query_total: Family<QueryLabels, Counter>,
    pub query_duration_seconds: Family<QueryLabels, Histogram>,
    registry: Mutex<Registry>,
}

#[cfg(feature = "distributed")]
impl PartitionNodeMetrics {
    /// Create a new metrics registry with pre-registered metrics.
    pub fn new() -> Arc<Self> {
        let mut registry = Registry::default();

        let query_total: Family<QueryLabels, Counter> = Family::default();
        let query_duration_seconds: Family<QueryLabels, Histogram> = Family::new_with_constructor(|| {
            Histogram::new(prometheus_client::metrics::histogram::exponential_buckets(
                0.0001, 2.0, 12,
            ))
        });

        registry.register(
            "partition_query",
            "Total number of partition queries",
            query_total.clone(),
        );
        registry.register(
            "partition_query_duration_seconds",
            "Query duration in seconds",
            query_duration_seconds.clone(),
        );

        Arc::new(Self {
            query_total,
            query_duration_seconds,
            registry: Mutex::new(registry),
        })
    }

    /// Encode the registry into Prometheus text format.
    pub fn encode(&self) -> String {
        let mut buf = String::new();
        let registry = self.registry.lock().expect("metrics registry lock poisoned");
        encode(&mut buf, &registry).expect("prometheus encode failed");
        buf
    }

    /// Record a query observation.
    pub fn record_query(&self, backend: &str, duration_secs: f64, ok: bool) {
        let labels = QueryLabels {
            backend: backend.to_string(),
            status: if ok { "ok" } else { "error" }.to_string(),
        };
        self.query_total.get_or_create(&labels).inc();
        self.query_duration_seconds
            .get_or_create(&labels)
            .observe(duration_secs);
    }
}
