use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use parking_lot::RwLock;

/// Metrics collector for system observability
pub struct MetricsCollector {
    /// Task execution count
    task_count: Arc<AtomicU64>,

    /// Task execution success count
    task_success_count: Arc<AtomicU64>,

    /// Task execution failure count
    task_failure_count: Arc<AtomicU64>,

    /// Total task execution time (milliseconds)
    task_duration_total: Arc<AtomicU64>,

    /// Agent heartbeat count
    heartbeat_count: Arc<AtomicU64>,

    /// Swarm coalition formations
    coalition_count: Arc<AtomicU64>,

    /// Provider API calls
    provider_calls: Arc<RwLock<HashMap<String, u64>>>,

    /// Token usage tracking
    tokens_used: Arc<RwLock<HashMap<String, TokenMetrics>>>,

    /// Errors by type
    errors: Arc<RwLock<HashMap<String, u64>>>,
}

/// Token usage metrics for a provider
#[derive(Debug, Clone)]
pub struct TokenMetrics {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub total_cost: f64,
    pub call_count: u64,
}

impl MetricsCollector {
    /// Create new metrics collector
    pub fn new() -> Self {
        Self {
            task_count: Arc::new(AtomicU64::new(0)),
            task_success_count: Arc::new(AtomicU64::new(0)),
            task_failure_count: Arc::new(AtomicU64::new(0)),
            task_duration_total: Arc::new(AtomicU64::new(0)),
            heartbeat_count: Arc::new(AtomicU64::new(0)),
            coalition_count: Arc::new(AtomicU64::new(0)),
            provider_calls: Arc::new(RwLock::new(HashMap::new())),
            tokens_used: Arc::new(RwLock::new(HashMap::new())),
            errors: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Record task execution start
    pub fn record_task_start(&self) {
        self.task_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Record successful task completion
    pub fn record_task_success(&self, duration_ms: u64) {
        self.task_success_count.fetch_add(1, Ordering::Relaxed);
        self.task_duration_total
            .fetch_add(duration_ms, Ordering::Relaxed);
    }

    /// Record task failure
    pub fn record_task_failure(&self, duration_ms: u64, error_type: &str) {
        self.task_failure_count.fetch_add(1, Ordering::Relaxed);
        self.task_duration_total
            .fetch_add(duration_ms, Ordering::Relaxed);

        let mut errors = self.errors.write();
        *errors.entry(error_type.to_string()).or_insert(0) += 1;
    }

    /// Record agent heartbeat
    pub fn record_heartbeat(&self) {
        self.heartbeat_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Record coalition formation
    pub fn record_coalition(&self) {
        self.coalition_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Record provider API call with token usage
    pub fn record_provider_call(
        &self,
        provider: &str,
        input_tokens: u64,
        output_tokens: u64,
        cost: f64,
    ) {
        // Increment call count
        let mut calls = self.provider_calls.write();
        *calls.entry(provider.to_string()).or_insert(0) += 1;
        drop(calls);

        // Track token usage
        let mut tokens = self.tokens_used.write();
        let metrics = tokens
            .entry(provider.to_string())
            .or_insert_with(|| TokenMetrics {
                input_tokens: 0,
                output_tokens: 0,
                total_cost: 0.0,
                call_count: 0,
            });

        metrics.input_tokens += input_tokens;
        metrics.output_tokens += output_tokens;
        metrics.total_cost += cost;
        metrics.call_count += 1;
    }

    /// Get current task metrics
    pub fn get_task_metrics(&self) -> TaskMetrics {
        let total = self.task_count.load(Ordering::Relaxed);
        let success = self.task_success_count.load(Ordering::Relaxed);
        let failure = self.task_failure_count.load(Ordering::Relaxed);
        let duration_total = self.task_duration_total.load(Ordering::Relaxed);

        let success_rate = if total > 0 {
            (success as f64 / total as f64) * 100.0
        } else {
            0.0
        };

        let avg_duration = if total > 0 {
            duration_total as f64 / total as f64
        } else {
            0.0
        };

        TaskMetrics {
            total_tasks: total,
            successful_tasks: success,
            failed_tasks: failure,
            success_rate,
            avg_duration_ms: avg_duration,
            total_duration_ms: duration_total,
        }
    }

    /// Get provider metrics
    pub fn get_provider_metrics(&self) -> HashMap<String, ProviderMetrics> {
        let calls = self.provider_calls.read();
        let tokens = self.tokens_used.read();

        let mut result = HashMap::new();

        for (provider, _) in calls.iter() {
            let call_count = calls.get(provider).copied().unwrap_or(0);
            let token_data = tokens.get(provider).cloned();

            let input_tokens = token_data.as_ref().map(|t| t.input_tokens).unwrap_or(0);
            let output_tokens = token_data.as_ref().map(|t| t.output_tokens).unwrap_or(0);
            let total_cost = token_data.as_ref().map(|t| t.total_cost).unwrap_or(0.0);

            result.insert(
                provider.clone(),
                ProviderMetrics {
                    call_count,
                    input_tokens,
                    output_tokens,
                    total_cost,
                },
            );
        }

        result
    }

    /// Get system-wide metrics
    pub fn get_system_metrics(&self) -> SystemMetrics {
        let task_metrics = self.get_task_metrics();
        let heartbeats = self.heartbeat_count.load(Ordering::Relaxed);
        let coalitions = self.coalition_count.load(Ordering::Relaxed);
        let error_map = self.errors.read().clone();

        let mut error_counts = Vec::new();
        for (error_type, count) in error_map.iter() {
            error_counts.push((error_type.clone(), *count));
        }
        error_counts.sort_by(|a, b| b.1.cmp(&a.1));

        SystemMetrics {
            tasks: task_metrics,
            heartbeats,
            coalitions_formed: coalitions,
            top_errors: error_counts.into_iter().take(10).collect(),
        }
    }

    /// Get error distribution
    pub fn get_errors(&self) -> HashMap<String, u64> {
        self.errors.read().clone()
    }

    /// Reset all metrics (for testing)
    #[cfg(test)]
    pub fn reset(&self) {
        self.task_count.store(0, Ordering::Relaxed);
        self.task_success_count.store(0, Ordering::Relaxed);
        self.task_failure_count.store(0, Ordering::Relaxed);
        self.task_duration_total.store(0, Ordering::Relaxed);
        self.heartbeat_count.store(0, Ordering::Relaxed);
        self.coalition_count.store(0, Ordering::Relaxed);
        self.provider_calls.write().clear();
        self.tokens_used.write().clear();
        self.errors.write().clear();
    }
}

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

impl Clone for MetricsCollector {
    fn clone(&self) -> Self {
        Self {
            task_count: Arc::clone(&self.task_count),
            task_success_count: Arc::clone(&self.task_success_count),
            task_failure_count: Arc::clone(&self.task_failure_count),
            task_duration_total: Arc::clone(&self.task_duration_total),
            heartbeat_count: Arc::clone(&self.heartbeat_count),
            coalition_count: Arc::clone(&self.coalition_count),
            provider_calls: Arc::clone(&self.provider_calls),
            tokens_used: Arc::clone(&self.tokens_used),
            errors: Arc::clone(&self.errors),
        }
    }
}

/// Task execution metrics
#[derive(Debug, Clone)]
pub struct TaskMetrics {
    pub total_tasks: u64,
    pub successful_tasks: u64,
    pub failed_tasks: u64,
    pub success_rate: f64,
    pub avg_duration_ms: f64,
    pub total_duration_ms: u64,
}

/// Provider-specific metrics
#[derive(Debug, Clone)]
pub struct ProviderMetrics {
    pub call_count: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub total_cost: f64,
}

/// System-wide metrics
#[derive(Debug, Clone)]
pub struct SystemMetrics {
    pub tasks: TaskMetrics,
    pub heartbeats: u64,
    pub coalitions_formed: u64,
    pub top_errors: Vec<(String, u64)>,
}

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

    #[test]
    fn test_task_metrics() {
        let collector = MetricsCollector::new();

        collector.record_task_start();
        collector.record_task_start();
        collector.record_task_success(100);
        collector.record_task_failure(200, "timeout");

        let metrics = collector.get_task_metrics();
        assert_eq!(metrics.total_tasks, 2);
        assert_eq!(metrics.successful_tasks, 1);
        assert_eq!(metrics.failed_tasks, 1);
        assert_eq!(metrics.success_rate, 50.0);
    }

    #[test]
    fn test_provider_metrics() {
        let collector = MetricsCollector::new();

        collector.record_provider_call("claude", 100, 50, 0.5);
        collector.record_provider_call("claude", 200, 100, 1.0);
        collector.record_provider_call("openai", 150, 75, 0.75);

        let metrics = collector.get_provider_metrics();
        assert_eq!(metrics.get("claude").map(|m| m.call_count), Some(2));
        assert_eq!(metrics.get("openai").map(|m| m.call_count), Some(1));
    }

    #[test]
    fn test_heartbeat_tracking() {
        let collector = MetricsCollector::new();

        for _ in 0..10 {
            collector.record_heartbeat();
        }

        let system_metrics = collector.get_system_metrics();
        assert_eq!(system_metrics.heartbeats, 10);
    }

    #[test]
    fn test_coalition_tracking() {
        let collector = MetricsCollector::new();

        for _ in 0..5 {
            collector.record_coalition();
        }

        let system_metrics = collector.get_system_metrics();
        assert_eq!(system_metrics.coalitions_formed, 5);
    }

    #[test]
    fn test_error_tracking() {
        let collector = MetricsCollector::new();

        collector.record_task_failure(100, "timeout");
        collector.record_task_failure(200, "timeout");
        collector.record_task_failure(150, "network_error");

        let errors = collector.get_errors();
        assert_eq!(errors.get("timeout"), Some(&2));
        assert_eq!(errors.get("network_error"), Some(&1));
    }

    #[test]
    fn test_system_metrics() {
        let collector = MetricsCollector::new();

        collector.record_task_start();
        collector.record_task_success(100);
        collector.record_heartbeat();
        collector.record_coalition();

        let metrics = collector.get_system_metrics();
        assert_eq!(metrics.tasks.total_tasks, 1);
        assert_eq!(metrics.heartbeats, 1);
        assert_eq!(metrics.coalitions_formed, 1);
    }

    #[test]
    fn test_metrics_cloning() {
        let collector1 = MetricsCollector::new();
        collector1.record_task_start();
        collector1.record_task_success(100);

        let collector2 = collector1.clone();
        let metrics2 = collector2.get_task_metrics();
        assert_eq!(metrics2.total_tasks, 1);
    }
}
