//! Usage statistics tracking for the gateway
//!
//! Tracks invocations, cache hits, tools discovered, and calculates token/cost savings.

use std::sync::atomic::{AtomicU64, Ordering};

use dashmap::DashMap;
use serde::{Deserialize, Serialize};

/// Usage statistics for the gateway
pub struct UsageStats {
    /// Total tool invocations via `gateway_invoke`
    total_invocations: AtomicU64,
    /// Cache hits from response cache
    cache_hits: AtomicU64,
    /// Tools discovered via `gateway_search_tools`
    tools_discovered: AtomicU64,
    /// Per-tool usage counts (key = "server:tool")
    tool_usage: DashMap<String, AtomicU64>,
}

impl UsageStats {
    /// Create new statistics tracker
    #[must_use]
    pub fn new() -> Self {
        Self {
            total_invocations: AtomicU64::new(0),
            cache_hits: AtomicU64::new(0),
            tools_discovered: AtomicU64::new(0),
            tool_usage: DashMap::new(),
        }
    }

    /// Record a tool invocation
    pub fn record_invocation(&self, server: &str, tool: &str) {
        self.total_invocations.fetch_add(1, Ordering::Relaxed);
        let key = format!("{server}:{tool}");
        self.tool_usage
            .entry(key)
            .or_insert_with(|| AtomicU64::new(0))
            .fetch_add(1, Ordering::Relaxed);
    }

    /// Record a cache hit
    pub fn record_cache_hit(&self) {
        self.cache_hits.fetch_add(1, Ordering::Relaxed);
    }

    /// Record tools discovered in a search
    pub fn record_search(&self, count: u64) {
        self.tools_discovered.fetch_add(count, Ordering::Relaxed);
    }

    /// Get usage count for a specific tool
    pub fn tool_usage(&self, server: &str, tool: &str) -> u64 {
        let key = format!("{server}:{tool}");
        self.tool_usage
            .get(&key)
            .map_or(0, |entry| entry.load(Ordering::Relaxed))
    }

    /// Get snapshot of current statistics
    pub fn snapshot(&self, total_backend_tools: usize) -> StatsSnapshot {
        let invocations = self.total_invocations.load(Ordering::Relaxed);
        let cache_hits = self.cache_hits.load(Ordering::Relaxed);
        let discovered = self.tools_discovered.load(Ordering::Relaxed);

        // Calculate token savings
        // Without gateway: each invocation would load ALL backend tools (~150 tokens each)
        // With gateway: 4 meta-tools are loaded instead
        // Savings = (total_backend_tools - 4) * 150 tokens * invocations
        let tokens_saved = if total_backend_tools > 4 {
            (total_backend_tools - 4) as u64 * 150 * invocations
        } else {
            0
        };

        // Get top tools
        let mut tool_counts: Vec<(String, u64)> = self
            .tool_usage
            .iter()
            .map(|entry| (entry.key().clone(), entry.value().load(Ordering::Relaxed)))
            .collect();
        tool_counts.sort_by(|a, b| b.1.cmp(&a.1));
        tool_counts.truncate(10);

        let top_tools: Vec<TopTool> = tool_counts
            .into_iter()
            .map(|(name, count)| {
                let parts: Vec<&str> = name.split(':').collect();
                TopTool {
                    server: parts.first().unwrap_or(&"").to_string(),
                    tool: parts.get(1).unwrap_or(&"").to_string(),
                    count,
                }
            })
            .collect();

        #[allow(clippy::cast_precision_loss)]
        let cache_hit_rate = if invocations > 0 {
            cache_hits as f64 / invocations as f64
        } else {
            0.0
        };

        StatsSnapshot {
            invocations,
            cache_hits,
            cache_hit_rate,
            tools_discovered: discovered,
            tools_available: total_backend_tools,
            tokens_saved,
            top_tools,
        }
    }

    /// Calculate estimated cost savings
    #[allow(clippy::cast_precision_loss)]
    pub fn cost_savings(&self, total_backend_tools: usize, price_per_million: f64) -> f64 {
        let snapshot = self.snapshot(total_backend_tools);
        snapshot.tokens_saved as f64 * price_per_million / 1_000_000.0
    }
}

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

/// Snapshot of usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StatsSnapshot {
    /// Total invocations
    pub invocations: u64,
    /// Cache hits
    pub cache_hits: u64,
    /// Cache hit rate (0.0-1.0)
    pub cache_hit_rate: f64,
    /// Tools discovered via search
    pub tools_discovered: u64,
    /// Total tools available across backends
    pub tools_available: usize,
    /// Estimated tokens saved by using gateway
    pub tokens_saved: u64,
    /// Top 10 most-used tools
    pub top_tools: Vec<TopTool>,
}

impl StatsSnapshot {
    /// Calculate estimated cost savings at a given token price
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn estimated_savings_usd(&self, price_per_million: f64) -> f64 {
        self.tokens_saved as f64 * price_per_million / 1_000_000.0
    }
}

/// Top tool usage entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopTool {
    /// Server name
    pub server: String,
    /// Tool name
    pub tool: String,
    /// Usage count
    pub count: u64,
}

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

    #[test]
    fn test_record_invocation() {
        let stats = UsageStats::new();
        stats.record_invocation("server1", "tool1");
        stats.record_invocation("server1", "tool1");
        stats.record_invocation("server2", "tool2");

        assert_eq!(stats.tool_usage("server1", "tool1"), 2);
        assert_eq!(stats.tool_usage("server2", "tool2"), 1);
        assert_eq!(stats.tool_usage("server3", "tool3"), 0);
    }

    #[test]
    fn test_snapshot() {
        let stats = UsageStats::new();
        stats.record_invocation("server1", "tool1");
        stats.record_invocation("server1", "tool1");
        stats.record_invocation("server2", "tool2");
        stats.record_cache_hit();
        stats.record_search(5);

        let snapshot = stats.snapshot(100);

        assert_eq!(snapshot.invocations, 3);
        assert_eq!(snapshot.cache_hits, 1);
        assert!((snapshot.cache_hit_rate - 0.333).abs() < 0.01);
        assert_eq!(snapshot.tools_discovered, 5);
        assert_eq!(snapshot.tools_available, 100);
        // (100 - 4) * 150 * 3 = 43,200
        assert_eq!(snapshot.tokens_saved, 43_200);
    }

    #[test]
    fn test_cost_savings() {
        let stats = UsageStats::new();
        for _ in 0..100 {
            stats.record_invocation("server1", "tool1");
        }

        // Price: $15/million input tokens (Claude Opus 4.6)
        let savings = stats.cost_savings(100, 15.0);

        // (100 - 4) * 150 * 100 = 1,440,000 tokens
        // 1,440,000 * $15 / 1,000,000 = $21.60
        assert!((savings - 21.6).abs() < 0.01);
    }

    #[test]
    fn test_top_tools() {
        let stats = UsageStats::new();
        stats.record_invocation("s1", "popular");
        stats.record_invocation("s1", "popular");
        stats.record_invocation("s1", "popular");
        stats.record_invocation("s2", "rare");

        let snapshot = stats.snapshot(50);

        assert_eq!(snapshot.top_tools.len(), 2);
        assert_eq!(snapshot.top_tools[0].tool, "popular");
        assert_eq!(snapshot.top_tools[0].count, 3);
        assert_eq!(snapshot.top_tools[1].tool, "rare");
        assert_eq!(snapshot.top_tools[1].count, 1);
    }

    #[test]
    fn test_no_savings_with_few_tools() {
        let stats = UsageStats::new();
        stats.record_invocation("s1", "t1");

        // Only 3 tools available, gateway has 4 meta-tools
        let snapshot = stats.snapshot(3);
        assert_eq!(snapshot.tokens_saved, 0);
    }

    #[test]
    fn test_default_impl() {
        let stats = UsageStats::default();
        let snapshot = stats.snapshot(100);
        assert_eq!(snapshot.invocations, 0);
        assert_eq!(snapshot.cache_hits, 0);
    }

    #[test]
    fn test_cache_hit_tracking() {
        let stats = UsageStats::new();
        stats.record_invocation("s1", "t1");
        stats.record_invocation("s1", "t1");
        stats.record_cache_hit();

        let snapshot = stats.snapshot(50);
        assert_eq!(snapshot.invocations, 2);
        assert_eq!(snapshot.cache_hits, 1);
        assert!((snapshot.cache_hit_rate - 0.5).abs() < 0.01);
    }

    #[test]
    fn test_search_tracking() {
        let stats = UsageStats::new();
        stats.record_search(10);
        stats.record_search(5);

        let snapshot = stats.snapshot(100);
        assert_eq!(snapshot.tools_discovered, 15);
    }

    #[test]
    fn test_snapshot_estimated_savings() {
        let stats = UsageStats::new();
        stats.record_invocation("s1", "t1");

        let snapshot = stats.snapshot(100);
        let savings = snapshot.estimated_savings_usd(15.0);

        // (100 - 4) * 150 * 1 = 14,400 tokens
        // 14,400 * $15 / 1,000,000 = $0.216
        assert!((savings - 0.216).abs() < 0.001);
    }

    #[test]
    fn test_zero_invocations_cache_rate() {
        let stats = UsageStats::new();
        let snapshot = stats.snapshot(50);
        assert_eq!(snapshot.cache_hit_rate, 0.0);
    }

    #[test]
    fn test_top_tools_sorting() {
        let stats = UsageStats::new();
        stats.record_invocation("s1", "rare");
        stats.record_invocation("s2", "common");
        stats.record_invocation("s2", "common");
        stats.record_invocation("s2", "common");
        stats.record_invocation("s3", "medium");
        stats.record_invocation("s3", "medium");

        let snapshot = stats.snapshot(50);

        assert_eq!(snapshot.top_tools.len(), 3);
        assert_eq!(snapshot.top_tools[0].tool, "common");
        assert_eq!(snapshot.top_tools[0].count, 3);
        assert_eq!(snapshot.top_tools[1].tool, "medium");
        assert_eq!(snapshot.top_tools[1].count, 2);
        assert_eq!(snapshot.top_tools[2].tool, "rare");
        assert_eq!(snapshot.top_tools[2].count, 1);
    }
}
