//! Telemetry aggregation at the relay edge.
//!
//! Collects per-connection and per-stream metrics and emits them
//! to the control plane metrics ingest pipeline.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::info;

/// Per-connection telemetry snapshot.
#[derive(Debug, Clone, Default)]
pub struct ConnectionMetrics {
    pub session_id: String,
    pub rtt_us: u64,
    pub packet_loss_pct: f32,
    pub bytes_sent: u64,
    pub bytes_received: u64,
    pub active_streams: u32,
    pub congestion_state: String,
    pub relay_path: Vec<String>,
    pub jitter_ms: f32,
    pub fec_recovered: u32,
}

/// Aggregated relay-level telemetry.
#[derive(Debug, Clone, Default)]
pub struct RelayMetrics {
    pub total_connections: u32,
    pub total_bandwidth_bps: u64,
    pub avg_rtt_us: u64,
    pub total_packets_forwarded: u64,
    pub total_packets_dropped: u64,
}

/// Telemetry collector that aggregates per-connection metrics.
pub struct TelemetryCollector {
    connections: Arc<RwLock<HashMap<String, ConnectionMetrics>>>,
    total_forwarded: Arc<RwLock<u64>>,
    total_dropped: Arc<RwLock<u64>>,
}

impl TelemetryCollector {
    pub fn new() -> Self {
        Self {
            connections: Arc::new(RwLock::new(HashMap::new())),
            total_forwarded: Arc::new(RwLock::new(0)),
            total_dropped: Arc::new(RwLock::new(0)),
        }
    }

    /// Update metrics for a specific connection.
    pub async fn update(&self, session_id: &str, metrics: ConnectionMetrics) {
        self.connections
            .write()
            .await
            .insert(session_id.to_string(), metrics);
    }

    /// Record packets forwarded / dropped.
    pub async fn record_forward(&self, count: u64) {
        *self.total_forwarded.write().await += count;
    }

    pub async fn record_drop(&self, count: u64) {
        *self.total_dropped.write().await += count;
    }

    /// Get aggregated relay-level metrics snapshot.
    pub async fn aggregate(&self) -> RelayMetrics {
        let connections = self.connections.read().await;
        let count = connections.len() as u32;

        let total_bandwidth: u64 = connections
            .values()
            .map(|c| c.bytes_sent + c.bytes_received)
            .sum();

        let avg_rtt = if count > 0 {
            connections.values().map(|c| c.rtt_us).sum::<u64>() / count as u64
        } else {
            0
        };

        RelayMetrics {
            total_connections: count,
            total_bandwidth_bps: total_bandwidth,
            avg_rtt_us: avg_rtt,
            total_packets_forwarded: *self.total_forwarded.read().await,
            total_packets_dropped: *self.total_dropped.read().await,
        }
    }

    /// Get metrics for a single connection.
    pub async fn get_connection(&self, session_id: &str) -> Option<ConnectionMetrics> {
        self.connections.read().await.get(session_id).cloned()
    }

    /// Remove a connection's metrics (on disconnect).
    pub async fn remove(&self, session_id: &str) {
        self.connections.write().await.remove(session_id);
    }
}

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

    fn conn(
        session_id: &str,
        rtt_us: u64,
        bytes_sent: u64,
        bytes_received: u64,
    ) -> ConnectionMetrics {
        ConnectionMetrics {
            session_id: session_id.to_string(),
            rtt_us,
            bytes_sent,
            bytes_received,
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn aggregate_empty_is_zeroed() {
        let collector = TelemetryCollector::new();
        let agg = collector.aggregate().await;
        assert_eq!(agg.total_connections, 0);
        assert_eq!(agg.total_bandwidth_bps, 0);
        assert_eq!(agg.avg_rtt_us, 0);
        assert_eq!(agg.total_packets_forwarded, 0);
        assert_eq!(agg.total_packets_dropped, 0);
    }

    #[tokio::test]
    async fn aggregate_sums_and_averages_connections() {
        let collector = TelemetryCollector::new();

        collector.update("s1", conn("s1", 10_000, 100, 900)).await;
        collector
            .update("s2", conn("s2", 30_000, 1_000, 2_000))
            .await;

        collector.record_forward(7).await;
        collector.record_drop(2).await;

        let agg = collector.aggregate().await;
        assert_eq!(agg.total_connections, 2);
        assert_eq!(agg.total_bandwidth_bps, (100 + 900) + (1_000 + 2_000));
        assert_eq!(agg.avg_rtt_us, (10_000 + 30_000) / 2);
        assert_eq!(agg.total_packets_forwarded, 7);
        assert_eq!(agg.total_packets_dropped, 2);
    }

    #[tokio::test]
    async fn remove_connection_updates_aggregate() {
        let collector = TelemetryCollector::new();
        collector.update("s1", conn("s1", 10_000, 100, 200)).await;
        collector.update("s2", conn("s2", 30_000, 300, 400)).await;

        collector.remove("s1").await;
        let agg = collector.aggregate().await;
        assert_eq!(agg.total_connections, 1);
        assert_eq!(agg.avg_rtt_us, 30_000);
        assert_eq!(agg.total_bandwidth_bps, 700);

        assert!(collector.get_connection("s1").await.is_none());
        assert!(collector.get_connection("s2").await.is_some());
    }
}
