//! Federation Metrics - Lightweight metrics collection
//!
//! Provides Counter, Gauge, and Histogram metric types with label support.
//! Exports to Prometheus text format.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
use std::time::Instant;

/// Metric label set: `vec![("store", "social"), ("op", "query")]`
pub type Labels = Vec<(String, String)>;

/// A single counter metric (monotonically increasing)
#[derive(Debug, Default)]
struct Counter {
    value: AtomicU64,
}

/// A single gauge metric (can go up and down)
#[derive(Debug, Default)]
struct Gauge {
    value: AtomicU64,
}

/// A simple histogram backed by a Vec of observations
#[derive(Debug, Default)]
struct Histogram {
    observations: RwLock<Vec<f64>>,
}

/// Internal metric entry
enum MetricEntry {
    Counter(Counter),
    Gauge(Gauge),
    Histogram(Histogram),
}

/// Federation metrics collector
///
/// Thread-safe, lock-free counters/gauges with a simple histogram implementation.
pub struct FederationMetrics {
    /// name+labels → metric entry
    entries: RwLock<HashMap<String, MetricEntry>>,
}

impl FederationMetrics {
    pub fn new() -> Self {
        Self {
            entries: RwLock::new(HashMap::new()),
        }
    }

    // ========== Counter ==========

    /// Increment a counter by 1
    pub fn inc_counter(&self, name: &str, labels: &Labels) {
        let key = Self::make_key(name, labels);
        // Fast path: read lock
        {
            let entries = self.entries.read().unwrap();
            if let Some(MetricEntry::Counter(c)) = entries.get(&key) {
                c.value.fetch_add(1, Ordering::Relaxed);
                return;
            }
        }
        // Slow path: write lock to create
        let mut entries = self.entries.write().unwrap();
        let entry = entries
            .entry(key)
            .or_insert_with(|| MetricEntry::Counter(Counter::default()));
        if let MetricEntry::Counter(c) = entry {
            c.value.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Increment a counter by `n`
    pub fn add_counter(&self, name: &str, labels: &Labels, n: u64) {
        let key = Self::make_key(name, labels);
        {
            let entries = self.entries.read().unwrap();
            if let Some(MetricEntry::Counter(c)) = entries.get(&key) {
                c.value.fetch_add(n, Ordering::Relaxed);
                return;
            }
        }
        let mut entries = self.entries.write().unwrap();
        let entry = entries
            .entry(key)
            .or_insert_with(|| MetricEntry::Counter(Counter::default()));
        if let MetricEntry::Counter(c) = entry {
            c.value.fetch_add(n, Ordering::Relaxed);
        }
    }

    // ========== Gauge ==========

    /// Set a gauge value
    pub fn set_gauge(&self, name: &str, labels: &Labels, value: u64) {
        let key = Self::make_key(name, labels);
        {
            let entries = self.entries.read().unwrap();
            if let Some(MetricEntry::Gauge(g)) = entries.get(&key) {
                g.value.store(value, Ordering::Relaxed);
                return;
            }
        }
        let mut entries = self.entries.write().unwrap();
        let entry = entries
            .entry(key)
            .or_insert_with(|| MetricEntry::Gauge(Gauge::default()));
        if let MetricEntry::Gauge(g) = entry {
            g.value.store(value, Ordering::Relaxed);
        }
    }

    // ========== Histogram ==========

    /// Record an observation (e.g. latency in ms)
    pub fn observe(&self, name: &str, labels: &Labels, value: f64) {
        let key = Self::make_key(name, labels);
        {
            let entries = self.entries.read().unwrap();
            if let Some(MetricEntry::Histogram(h)) = entries.get(&key) {
                h.observations.write().unwrap().push(value);
                return;
            }
        }
        let mut entries = self.entries.write().unwrap();
        let entry = entries
            .entry(key)
            .or_insert_with(|| MetricEntry::Histogram(Histogram::default()));
        if let MetricEntry::Histogram(h) = entry {
            h.observations.write().unwrap().push(value);
        }
    }

    // ========== Timer helper ==========

    /// Start a timer; call `.finish()` on the returned guard to record duration.
    pub fn start_timer(&self, name: &str, labels: Labels) -> TimerGuard<'_> {
        TimerGuard {
            metrics: self,
            name: name.to_string(),
            labels,
            start: Instant::now(),
        }
    }

    // ========== Export ==========

    /// Export all metrics in Prometheus text exposition format
    pub fn export_prometheus(&self) -> String {
        let entries = self.entries.read().unwrap();
        let mut output = String::new();

        // Sort keys for stable output
        let mut keys: Vec<&String> = entries.keys().collect();
        keys.sort();

        for key in keys {
            if let Some(entry) = entries.get(key) {
                match entry {
                    MetricEntry::Counter(c) => {
                        let val = c.value.load(Ordering::Relaxed);
                        output.push_str(&format!("{} {}\n", key, val));
                    }
                    MetricEntry::Gauge(g) => {
                        let val = g.value.load(Ordering::Relaxed);
                        output.push_str(&format!("{} {}\n", key, val));
                    }
                    MetricEntry::Histogram(h) => {
                        let obs = h.observations.read().unwrap();
                        if obs.is_empty() {
                            output.push_str(&format!("{}_count 0\n", key));
                            continue;
                        }
                        let count = obs.len();
                        let sum: f64 = obs.iter().sum();
                        let mut sorted = obs.clone();
                        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                        let p50 = sorted.get(count / 2).copied().unwrap_or(0.0);
                        let p99_idx = ((count as f64) * 0.99) as usize;
                        let p99 = sorted
                            .get(p99_idx.min(count.saturating_sub(1)))
                            .copied()
                            .unwrap_or(0.0);

                        output.push_str(&format!("{}_count {}\n", key, count));
                        output.push_str(&format!("{}_sum {:.3}\n", key, sum));
                        output.push_str(&format!("{}_p50 {:.3}\n", key, p50));
                        output.push_str(&format!("{}_p99 {:.3}\n", key, p99));
                    }
                }
            }
        }

        output
    }

    /// Reset all metrics (useful for testing)
    pub fn reset(&self) {
        let mut entries = self.entries.write().unwrap();
        entries.clear();
    }

    // -- helpers --

    fn make_key(name: &str, labels: &Labels) -> String {
        if labels.is_empty() {
            return name.to_string();
        }
        let label_str: Vec<String> = labels.iter().map(|(k, v)| format!("{}=\"{}\"", k, v)).collect();
        format!("{}{{{}}}", name, label_str.join(","))
    }
}

impl Default for FederationMetrics {
    fn default() -> Self {
        Self::new()
    }
}

/// Timer guard — records duration to a histogram when dropped or explicitly finished
pub struct TimerGuard<'a> {
    metrics: &'a FederationMetrics,
    name: String,
    labels: Labels,
    start: Instant,
}

impl<'a> TimerGuard<'a> {
    /// Finish the timer and record the duration in milliseconds
    pub fn finish(self) -> f64 {
        let duration_ms = self.start.elapsed().as_secs_f64() * 1000.0;
        self.metrics.observe(&self.name, &self.labels, duration_ms);
        duration_ms
    }
}

impl<'a> Drop for TimerGuard<'a> {
    fn drop(&mut self) {
        let duration_ms = self.start.elapsed().as_secs_f64() * 1000.0;
        self.metrics.observe(&self.name, &self.labels, duration_ms);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_counter() {
        let m = FederationMetrics::new();
        let labels = vec![("store".to_string(), "social".to_string())];

        m.inc_counter("query_total", &labels);
        m.inc_counter("query_total", &labels);
        m.inc_counter("query_total", &labels);

        let output = m.export_prometheus();
        assert!(output.contains("query_total{store=\"social\"} 3"));
    }

    #[test]
    fn test_gauge() {
        let m = FederationMetrics::new();
        let labels = vec![];

        m.set_gauge("tx_active", &labels, 5);
        let output = m.export_prometheus();
        assert!(output.contains("tx_active 5"));

        m.set_gauge("tx_active", &labels, 3);
        let output = m.export_prometheus();
        assert!(output.contains("tx_active 3"));
    }

    #[test]
    fn test_histogram() {
        let m = FederationMetrics::new();
        let labels = vec![("op".to_string(), "query".to_string())];

        for i in 0..100 {
            m.observe("query_duration_ms", &labels, i as f64);
        }

        let output = m.export_prometheus();
        assert!(output.contains("query_duration_ms{op=\"query\"}_count 100"));
        assert!(output.contains("query_duration_ms{op=\"query\"}_sum"));
        assert!(output.contains("query_duration_ms{op=\"query\"}_p50"));
        assert!(output.contains("query_duration_ms{op=\"query\"}_p99"));
    }

    #[test]
    fn test_reset() {
        let m = FederationMetrics::new();
        m.inc_counter("test", &vec![]);
        m.reset();
        let output = m.export_prometheus();
        assert!(output.is_empty());
    }
}
