hikari-desktop/src-tauri/src/stats.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
use crate::achievements::{AchievementProgress, check_achievements};

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct UsageStats {
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_cost_usd: f64,
    pub session_input_tokens: u64,
    pub session_output_tokens: u64,
    pub session_cost_usd: f64,
    pub model: Option<String>,

    // New fields
    pub messages_exchanged: u64,
    pub session_messages_exchanged: u64,
    pub code_blocks_generated: u64,
    pub session_code_blocks_generated: u64,
    pub files_edited: u64,
    pub session_files_edited: u64,
    pub files_created: u64,
    pub session_files_created: u64,
    pub tools_usage: HashMap<String, u64>,
    pub session_tools_usage: HashMap<String, u64>,
    pub session_duration_seconds: u64,
    #[serde(skip)]
    pub session_start: Option<Instant>,

    // Achievement tracking
    #[serde(skip)]
    pub achievements: AchievementProgress,
}

impl UsageStats {
    pub fn new() -> Self {
        let mut stats = Self::default();
        stats.achievements.start_session();
        stats
    }

    pub fn add_usage(&mut self, input_tokens: u64, output_tokens: u64, model: &str) {
        self.total_input_tokens += input_tokens;
        self.total_output_tokens += output_tokens;
        self.session_input_tokens += input_tokens;
        self.session_output_tokens += output_tokens;

        let cost = calculate_cost(input_tokens, output_tokens, model);
        self.total_cost_usd += cost;
        self.session_cost_usd += cost;

        self.model = Some(model.to_string());
    }

    pub fn reset_session(&mut self) {
        self.session_input_tokens = 0;
        self.session_output_tokens = 0;
        self.session_cost_usd = 0.0;
        self.session_messages_exchanged = 0;
        self.session_code_blocks_generated = 0;
        self.session_files_edited = 0;
        self.session_files_created = 0;
        self.session_tools_usage.clear();
        self.session_duration_seconds = 0;
        self.session_start = Some(Instant::now());
        self.achievements.start_session();
    }

    pub fn increment_messages(&mut self) {
        self.messages_exchanged += 1;
        self.session_messages_exchanged += 1;
    }

    pub fn increment_code_blocks(&mut self) {
        self.code_blocks_generated += 1;
        self.session_code_blocks_generated += 1;
    }

    pub fn increment_files_edited(&mut self) {
        self.files_edited += 1;
        self.session_files_edited += 1;
    }

    pub fn increment_files_created(&mut self) {
        self.files_created += 1;
        self.session_files_created += 1;
    }

    pub fn increment_tool_usage(&mut self, tool_name: &str) {
        *self.tools_usage.entry(tool_name.to_string()).or_insert(0) += 1;
        *self.session_tools_usage.entry(tool_name.to_string()).or_insert(0) += 1;
    }

    pub fn get_session_duration(&mut self) -> u64 {
        // Only update if more than 1 second has passed to reduce calculations
        if let Some(start) = self.session_start {
            let elapsed = start.elapsed().as_secs();
            if elapsed > self.session_duration_seconds {
                self.session_duration_seconds = elapsed;
            }
        }
        self.session_duration_seconds
    }

    pub fn check_achievements(&mut self) -> Vec<crate::achievements::AchievementId> {
        let stats_copy = UsageStats {
            total_input_tokens: self.total_input_tokens,
            total_output_tokens: self.total_output_tokens,
            total_cost_usd: self.total_cost_usd,
            session_input_tokens: self.session_input_tokens,
            session_output_tokens: self.session_output_tokens,
            session_cost_usd: self.session_cost_usd,
            model: self.model.clone(),
            messages_exchanged: self.messages_exchanged,
            session_messages_exchanged: self.session_messages_exchanged,
            code_blocks_generated: self.code_blocks_generated,
            session_code_blocks_generated: self.session_code_blocks_generated,
            files_edited: self.files_edited,
            session_files_edited: self.session_files_edited,
            files_created: self.files_created,
            session_files_created: self.session_files_created,
            tools_usage: self.tools_usage.clone(),
            session_tools_usage: self.session_tools_usage.clone(),
            session_duration_seconds: self.session_duration_seconds,
            session_start: self.session_start,
            achievements: AchievementProgress::new(), // Dummy for copy
        };
        check_achievements(&stats_copy, &mut self.achievements)
    }
}

// Pricing as of January 2025
// https://www.anthropic.com/pricing
fn calculate_cost(input_tokens: u64, output_tokens: u64, model: &str) -> f64 {
    let (input_price_per_million, output_price_per_million) = match model {
        // Opus 4.5
        "claude-opus-4-5-20251101" => (15.0, 75.0),

        // Opus 4
        "claude-opus-4-20250514" => (15.0, 75.0),

        // Sonnet 4
        "claude-sonnet-4-20250514" => (3.0, 15.0),

        // Previous generation models
        "claude-3-5-sonnet-20241022" => (3.0, 15.0),
        "claude-3-5-sonnet-20240620" => (3.0, 15.0),
        "claude-3-5-haiku-20241022" => (1.0, 5.0),
        "claude-3-opus-20240229" => (15.0, 75.0),
        "claude-3-sonnet-20240229" => (3.0, 15.0),
        "claude-3-haiku-20240307" => (0.25, 1.25),

        // Default to Sonnet pricing if model unknown
        _ => (3.0, 15.0),
    };

    let input_cost = (input_tokens as f64 / 1_000_000.0) * input_price_per_million;
    let output_cost = (output_tokens as f64 / 1_000_000.0) * output_price_per_million;

    input_cost + output_cost
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StatsUpdateEvent {
    pub stats: UsageStats,
}

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

    #[test]
    fn test_cost_calculation_sonnet() {
        let cost = calculate_cost(1000, 2000, "claude-sonnet-4-20250514");
        // 1000 input * $3/M = $0.003
        // 2000 output * $15/M = $0.030
        // Total = $0.033
        assert!((cost - 0.033).abs() < 0.0001);
    }

    #[test]
    fn test_cost_calculation_opus() {
        let cost = calculate_cost(1000, 2000, "claude-opus-4-20250514");
        // 1000 input * $15/M = $0.015
        // 2000 output * $75/M = $0.150
        // Total = $0.165
        assert!((cost - 0.165).abs() < 0.0001);
    }

    #[test]
    fn test_usage_stats_accumulation() {
        let mut stats = UsageStats::new();
        stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");

        assert_eq!(stats.total_input_tokens, 1000);
        assert_eq!(stats.total_output_tokens, 2000);
        assert_eq!(stats.session_input_tokens, 1000);
        assert_eq!(stats.session_output_tokens, 2000);
        assert!((stats.total_cost_usd - 0.033).abs() < 0.0001);
    }

    #[test]
    fn test_session_reset() {
        let mut stats = UsageStats::new();
        stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
        stats.reset_session();

        assert_eq!(stats.total_input_tokens, 1000);
        assert_eq!(stats.total_output_tokens, 2000);
        assert_eq!(stats.session_input_tokens, 0);
        assert_eq!(stats.session_output_tokens, 0);
        assert_eq!(stats.session_cost_usd, 0.0);
        assert!(stats.total_cost_usd > 0.0);
    }
}