//! Custom user-defined metrics created from config and manipulated via RPC.

use std::collections::HashMap;
use std::sync::Arc;

use folk_api::{CounterVec, GaugeVec, HistogramVec, MetricsRegistry};

use crate::config::{CollectorDef, CollectorType};

/// Holds all custom metrics created from `[[metrics.collectors]]` config.
pub struct CustomMetrics {
    counters: HashMap<String, Arc<dyn CounterVec>>,
    gauges: HashMap<String, Arc<dyn GaugeVec>>,
    histograms: HashMap<String, Arc<dyn HistogramVec>>,
    /// label keys per metric name, for validation
    label_keys: HashMap<String, Vec<String>>,
}

impl CustomMetrics {
    /// Create custom metrics from collector definitions, registering them in the
    /// provided MetricsRegistry. The `prefix` is prepended to metric names
    /// (e.g. prefix="folk", name="app_total" → "folk_app_total").
    pub fn from_config(
        defs: &[CollectorDef],
        prefix: &str,
        registry: &dyn MetricsRegistry,
    ) -> Self {
        let mut counters = HashMap::new();
        let mut gauges = HashMap::new();
        let mut histograms = HashMap::new();
        let mut label_keys = HashMap::new();

        for def in defs {
            let full_name = if prefix.is_empty() {
                def.name.clone()
            } else {
                format!("{prefix}_{}", def.name)
            };
            let labels_ref: Vec<&str> = def.labels.iter().map(String::as_str).collect();

            match def.collector_type {
                CollectorType::Counter => {
                    let cv = registry.counter_vec(&full_name, &def.help, &labels_ref);
                    counters.insert(def.name.clone(), cv);
                }
                CollectorType::Gauge => {
                    let gv = registry.gauge_vec(&full_name, &def.help, &labels_ref);
                    gauges.insert(def.name.clone(), gv);
                }
                CollectorType::Histogram => {
                    let hv = registry.histogram_vec(&full_name, &def.help, &labels_ref);
                    histograms.insert(def.name.clone(), hv);
                }
            }
            label_keys.insert(def.name.clone(), def.labels.clone());
        }

        Self {
            counters,
            gauges,
            histograms,
            label_keys,
        }
    }

    /// Increment a counter by `value` (default 1). Returns error if metric not found or wrong type.
    pub fn increment(
        &self,
        name: &str,
        labels: &[&str],
        value: u64,
    ) -> Result<(), CustomMetricError> {
        self.validate_labels(name, labels)?;
        let cv = self
            .counters
            .get(name)
            .ok_or_else(|| CustomMetricError::NotFound {
                name: name.into(),
                expected: "counter",
            })?;
        cv.with_labels(labels).inc_by(value);
        Ok(())
    }

    /// Observe a value on a histogram. Returns error if metric not found or wrong type.
    pub fn observe(
        &self,
        name: &str,
        labels: &[&str],
        value: f64,
    ) -> Result<(), CustomMetricError> {
        self.validate_labels(name, labels)?;
        let hv = self
            .histograms
            .get(name)
            .ok_or_else(|| CustomMetricError::NotFound {
                name: name.into(),
                expected: "histogram",
            })?;
        hv.with_labels(labels).observe(value);
        Ok(())
    }

    /// Set a gauge value. Returns error if metric not found or wrong type.
    pub fn set(&self, name: &str, labels: &[&str], value: i64) -> Result<(), CustomMetricError> {
        self.validate_labels(name, labels)?;
        let gv = self
            .gauges
            .get(name)
            .ok_or_else(|| CustomMetricError::NotFound {
                name: name.into(),
                expected: "gauge",
            })?;
        gv.with_labels(labels).set(value);
        Ok(())
    }

    fn validate_labels(&self, name: &str, labels: &[&str]) -> Result<(), CustomMetricError> {
        let Some(expected) = self.label_keys.get(name) else {
            return Ok(()); // will fail on lookup anyway
        };
        if labels.len() != expected.len() {
            return Err(CustomMetricError::LabelMismatch {
                name: name.into(),
                expected: expected.len(),
                got: labels.len(),
            });
        }
        Ok(())
    }
}

#[derive(Debug)]
pub enum CustomMetricError {
    NotFound {
        name: String,
        expected: &'static str,
    },
    LabelMismatch {
        name: String,
        expected: usize,
        got: usize,
    },
}

impl std::fmt::Display for CustomMetricError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotFound { name, expected } => {
                write!(f, "custom {expected} '{name}' not found")
            }
            Self::LabelMismatch {
                name,
                expected,
                got,
            } => {
                write!(f, "metric '{name}': expected {expected} labels, got {got}")
            }
        }
    }
}

impl std::error::Error for CustomMetricError {}
