// vapora-llm-router: Cost tracking module
// Track LLM API costs and usage statistics

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageStats {
    pub provider: String,
    pub task_type: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cost_cents: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostReport {
    pub total_cost_cents: u32,
    pub total_tasks: u32,
    pub cost_by_provider: HashMap<String, u32>,
    pub cost_by_task_type: HashMap<String, u32>,
    pub tokens_by_provider: HashMap<String, TokenCount>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TokenCount {
    pub input: u64,
    pub output: u64,
}

/// Thread-safe cost tracker
#[derive(Debug, Clone)]
pub struct CostTracker {
    inner: Arc<RwLock<CostTrackerInner>>,
}

#[derive(Debug, Default)]
struct CostTrackerInner {
    total_cost_cents: u32,
    cost_by_provider: HashMap<String, u32>,
    cost_by_task_type: HashMap<String, u32>,
    tokens_by_provider: HashMap<String, TokenCount>,
    tasks_completed: u32,
}

impl CostTracker {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(CostTrackerInner::default())),
        }
    }

    /// Log usage of a provider for a task
    pub fn log_usage(
        &self,
        provider: &str,
        task_type: &str,
        input_tokens: u64,
        output_tokens: u64,
        cost_cents: u32,
    ) {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        inner.total_cost_cents += cost_cents;
        inner.tasks_completed += 1;

        *inner
            .cost_by_provider
            .entry(provider.to_string())
            .or_insert(0) += cost_cents;

        *inner
            .cost_by_task_type
            .entry(task_type.to_string())
            .or_insert(0) += cost_cents;

        let token_count = inner
            .tokens_by_provider
            .entry(provider.to_string())
            .or_default();
        token_count.input += input_tokens;
        token_count.output += output_tokens;
    }

    /// Get total cost in USD
    pub fn total_cost_usd(&self) -> f64 {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.total_cost_cents as f64 / 100.0
    }

    /// Get number of completed tasks
    pub fn tasks_completed(&self) -> u32 {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.tasks_completed
    }

    /// Generate cost report
    pub fn generate_report(&self) -> CostReport {
        let inner = self.inner.read().expect("Failed to acquire read lock");

        CostReport {
            total_cost_cents: inner.total_cost_cents,
            total_tasks: inner.tasks_completed,
            cost_by_provider: inner.cost_by_provider.clone(),
            cost_by_task_type: inner.cost_by_task_type.clone(),
            tokens_by_provider: inner.tokens_by_provider.clone(),
        }
    }

    /// Get formatted report as string
    pub fn format_report(&self) -> String {
        let report = self.generate_report();

        let mut output = String::new();
        output.push_str(&format!(
            "=== Cost Report ===\nTotal Cost: ${:.2}\nTotal Tasks: {}\n\n",
            report.total_cost_cents as f64 / 100.0,
            report.total_tasks
        ));

        if !report.cost_by_provider.is_empty() {
            output.push_str("Cost by Provider:\n");
            for (provider, cost) in &report.cost_by_provider {
                output.push_str(&format!("  {}: ${:.2}\n", provider, *cost as f64 / 100.0));
            }
            output.push('\n');
        }

        if !report.cost_by_task_type.is_empty() {
            output.push_str("Cost by Task Type:\n");
            for (task_type, cost) in &report.cost_by_task_type {
                output.push_str(&format!("  {}: ${:.2}\n", task_type, *cost as f64 / 100.0));
            }
            output.push('\n');
        }

        if !report.tokens_by_provider.is_empty() {
            output.push_str("Tokens by Provider:\n");
            for (provider, tokens) in &report.tokens_by_provider {
                output.push_str(&format!(
                    "  {}: {} input, {} output\n",
                    provider, tokens.input, tokens.output
                ));
            }
        }

        output
    }

    /// Check if cost exceeds threshold
    pub fn exceeds_threshold(&self, threshold_cents: u32) -> bool {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.total_cost_cents > threshold_cents
    }

    /// Reset all statistics
    pub fn reset(&self) {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");
        inner.total_cost_cents = 0;
        inner.tasks_completed = 0;
        inner.cost_by_provider.clear();
        inner.cost_by_task_type.clear();
        inner.tokens_by_provider.clear();
    }
}

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

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

    #[test]
    fn test_cost_tracker() {
        let tracker = CostTracker::new();

        tracker.log_usage("claude", "development", 1000, 500, 10);
        tracker.log_usage("openai", "documentation", 800, 400, 8);
        tracker.log_usage("claude", "development", 1200, 600, 12);

        assert_eq!(tracker.total_cost_usd(), 0.30);
        assert_eq!(tracker.tasks_completed(), 3);

        let report = tracker.generate_report();
        assert_eq!(report.total_cost_cents, 30);
        assert_eq!(report.total_tasks, 3);
        assert_eq!(*report.cost_by_provider.get("claude").unwrap(), 22);
        assert_eq!(*report.cost_by_provider.get("openai").unwrap(), 8);
    }

    #[test]
    fn test_threshold_checking() {
        let tracker = CostTracker::new();
        tracker.log_usage("claude", "test", 1000, 500, 100);

        assert!(tracker.exceeds_threshold(50));
        assert!(!tracker.exceeds_threshold(200));
    }

    #[test]
    fn test_reset() {
        let tracker = CostTracker::new();
        tracker.log_usage("claude", "test", 1000, 500, 10);

        assert_eq!(tracker.tasks_completed(), 1);

        tracker.reset();

        assert_eq!(tracker.tasks_completed(), 0);
        assert_eq!(tracker.total_cost_usd(), 0.0);
    }
}
