//! Per-partition metrics snapshots used by the rebalancer.
//!
//! These are lightweight in-process counters. For Prometheus export,
//! a future integration layer would scrape these via an HTTP `/metrics`
//! endpoint using the `prometheus` crate.

#[cfg(feature = "distributed")]
use std::collections::HashMap;
#[cfg(feature = "distributed")]
use std::sync::Arc;
#[cfg(feature = "distributed")]
use std::time::{Duration, Instant};

#[cfg(feature = "distributed")]
use tokio::sync::RwLock;

// ── Snapshot ──────────────────────────────────────────────────────────────────

/// A point-in-time snapshot of a single partition's metrics.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct PartitionMetricsSnapshot {
    pub partition_id:    u32,
    pub node_count:      u64,
    pub edge_count:      u64,
    /// Queries processed in the last measurement window.
    pub query_count:     u64,
    pub sampled_at:      Instant,
}

#[cfg(feature = "distributed")]
impl PartitionMetricsSnapshot {
    /// Approximate bytes used (rough estimate: 256 bytes per node/edge).
    pub fn estimated_bytes(&self) -> u64 {
        (self.node_count + self.edge_count) * 256
    }
}

// ── Collector ─────────────────────────────────────────────────────────────────

/// Thread-safe metrics collector for all partitions on this pod.
#[cfg(feature = "distributed")]
#[derive(Clone)]
pub struct MetricsCollector {
    inner: Arc<RwLock<MetricsInner>>,
}

#[cfg(feature = "distributed")]
struct MetricsInner {
    partitions: HashMap<u32, PartitionMetricsSnapshot>,
}

#[cfg(feature = "distributed")]
impl MetricsCollector {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(MetricsInner {
                partitions: HashMap::new(),
            })),
        }
    }

    /// Record or update metrics for a partition.
    pub async fn record(
        &self,
        partition_id: u32,
        node_count:   u64,
        edge_count:   u64,
        query_count:  u64,
    ) {
        let mut inner = self.inner.write().await;
        inner.partitions.insert(partition_id, PartitionMetricsSnapshot {
            partition_id,
            node_count,
            edge_count,
            query_count,
            sampled_at: Instant::now(),
        });
    }

    /// Increment the query counter for a partition.
    pub async fn increment_query_count(&self, partition_id: u32) {
        let mut inner = self.inner.write().await;
        if let Some(snap) = inner.partitions.get_mut(&partition_id) {
            snap.query_count += 1;
        }
    }

    /// Return snapshots for all partitions.
    pub async fn snapshots(&self) -> Vec<PartitionMetricsSnapshot> {
        self.inner.read().await
            .partitions.values().cloned().collect()
    }

    /// Return the snapshot for a specific partition.
    pub async fn snapshot(&self, partition_id: u32) -> Option<PartitionMetricsSnapshot> {
        self.inner.read().await
            .partitions.get(&partition_id).cloned()
    }

    /// Return the age of the most recent sample for a partition.
    pub async fn snapshot_age(&self, partition_id: u32) -> Option<Duration> {
        self.inner.read().await
            .partitions.get(&partition_id)
            .map(|s| s.sampled_at.elapsed())
    }
}

#[cfg(feature = "distributed")]
impl Default for MetricsCollector {
    fn default() -> Self { Self::new() }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_record_and_retrieve() {
        let collector = MetricsCollector::new();
        collector.record(0, 100, 200, 50).await;
        let snap = collector.snapshot(0).await.unwrap();
        assert_eq!(snap.node_count, 100);
        assert_eq!(snap.edge_count, 200);
        assert_eq!(snap.query_count, 50);
    }

    #[tokio::test]
    async fn test_estimated_bytes() {
        let snap = PartitionMetricsSnapshot {
            partition_id: 0,
            node_count:   1000,
            edge_count:   2000,
            query_count:  0,
            sampled_at:   Instant::now(),
        };
        assert_eq!(snap.estimated_bytes(), 3000 * 256);
    }

    #[tokio::test]
    async fn test_increment_query_count() {
        let collector = MetricsCollector::new();
        collector.record(1, 0, 0, 0).await;
        collector.increment_query_count(1).await;
        collector.increment_query_count(1).await;
        let snap = collector.snapshot(1).await.unwrap();
        assert_eq!(snap.query_count, 2);
    }

    #[tokio::test]
    async fn test_snapshots_all() {
        let collector = MetricsCollector::new();
        collector.record(0, 10, 20, 5).await;
        collector.record(1, 30, 40, 10).await;
        let snaps = collector.snapshots().await;
        assert_eq!(snaps.len(), 2);
    }

    #[tokio::test]
    async fn test_missing_partition_returns_none() {
        let collector = MetricsCollector::new();
        assert!(collector.snapshot(99).await.is_none());
        assert!(collector.snapshot_age(99).await.is_none());
    }
}
