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

use crate::achievements::{check_achievements, AchievementProgress};
use chrono::{Local, Timelike};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
use tauri_plugin_store::StoreExt;

#[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>,

    // Extended tracking for achievements
    pub sessions_started: u64,
    pub consecutive_days: u64,
    pub total_days_used: u64,
    pub morning_sessions: u64, // Sessions started before 9 AM
    pub night_sessions: u64,   // Sessions started after 10 PM
    pub last_session_date: Option<String>, // ISO date string for streak tracking

    // 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();

        // Track session start for achievements
        self.track_session_start();
    }

    pub fn track_session_start(&mut self) {
        let now = Local::now();
        let today = now.format("%Y-%m-%d").to_string();
        let hour = now.hour();

        // Increment session count
        self.sessions_started += 1;

        // Track morning/night sessions
        if hour < 9 {
            self.morning_sessions += 1;
        }
        if hour >= 22 {
            self.night_sessions += 1;
        }

        // Track consecutive days and total days
        if let Some(last_date) = &self.last_session_date {
            if last_date != &today {
                // Check if it's the next day (consecutive)
                if is_consecutive_day(last_date, &today) {
                    self.consecutive_days += 1;
                } else {
                    // Streak broken
                    self.consecutive_days = 1;
                }
                self.total_days_used += 1;
                self.last_session_date = Some(today);
            }
            // Same day - don't increment anything
        } else {
            // First session ever
            self.consecutive_days = 1;
            self.total_days_used = 1;
            self.last_session_date = Some(today);
        }
    }

    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,
            sessions_started: self.sessions_started,
            consecutive_days: self.consecutive_days,
            total_days_used: self.total_days_used,
            morning_sessions: self.morning_sessions,
            night_sessions: self.night_sessions,
            last_session_date: self.last_session_date.clone(),
            achievements: AchievementProgress::new(), // Dummy for copy
        };
        check_achievements(&stats_copy, &mut self.achievements)
    }
}

// Helper function to check if two dates are consecutive
fn is_consecutive_day(prev_date: &str, current_date: &str) -> bool {
    use chrono::NaiveDate;

    let prev = NaiveDate::parse_from_str(prev_date, "%Y-%m-%d").ok();
    let current = NaiveDate::parse_from_str(current_date, "%Y-%m-%d").ok();

    match (prev, current) {
        (Some(p), Some(c)) => {
            let diff = c.signed_duration_since(p).num_days();
            diff == 1
        }
        _ => false,
    }
}

// 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,
}

/// Serializable struct for persisting only lifetime (total) stats
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PersistedStats {
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_cost_usd: f64,
    pub messages_exchanged: u64,
    pub code_blocks_generated: u64,
    pub files_edited: u64,
    pub files_created: u64,
    pub tools_usage: HashMap<String, u64>,
    pub sessions_started: u64,
    pub consecutive_days: u64,
    pub total_days_used: u64,
    pub morning_sessions: u64,
    pub night_sessions: u64,
    pub last_session_date: Option<String>,
}

impl From<&UsageStats> for PersistedStats {
    fn from(stats: &UsageStats) -> Self {
        PersistedStats {
            total_input_tokens: stats.total_input_tokens,
            total_output_tokens: stats.total_output_tokens,
            total_cost_usd: stats.total_cost_usd,
            messages_exchanged: stats.messages_exchanged,
            code_blocks_generated: stats.code_blocks_generated,
            files_edited: stats.files_edited,
            files_created: stats.files_created,
            tools_usage: stats.tools_usage.clone(),
            sessions_started: stats.sessions_started,
            consecutive_days: stats.consecutive_days,
            total_days_used: stats.total_days_used,
            morning_sessions: stats.morning_sessions,
            night_sessions: stats.night_sessions,
            last_session_date: stats.last_session_date.clone(),
        }
    }
}

impl UsageStats {
    /// Apply persisted stats to restore lifetime totals
    pub fn apply_persisted(&mut self, persisted: PersistedStats) {
        self.total_input_tokens = persisted.total_input_tokens;
        self.total_output_tokens = persisted.total_output_tokens;
        self.total_cost_usd = persisted.total_cost_usd;
        self.messages_exchanged = persisted.messages_exchanged;
        self.code_blocks_generated = persisted.code_blocks_generated;
        self.files_edited = persisted.files_edited;
        self.files_created = persisted.files_created;
        self.tools_usage = persisted.tools_usage;
        self.sessions_started = persisted.sessions_started;
        self.consecutive_days = persisted.consecutive_days;
        self.total_days_used = persisted.total_days_used;
        self.morning_sessions = persisted.morning_sessions;
        self.night_sessions = persisted.night_sessions;
        self.last_session_date = persisted.last_session_date;
    }
}

/// Save lifetime stats to persistent store
pub async fn save_stats(app: &tauri::AppHandle, stats: &UsageStats) -> Result<(), String> {
    let store = app.store("stats.json").map_err(|e| e.to_string())?;

    let persisted = PersistedStats::from(stats);

    println!("Saving stats: {:?}", persisted);

    store.set(
        "lifetime_stats",
        serde_json::to_value(persisted).map_err(|e| e.to_string())?,
    );
    store.save().map_err(|e| e.to_string())?;

    println!("Stats saved successfully");
    Ok(())
}

/// Load lifetime stats from persistent store
pub async fn load_stats(app: &tauri::AppHandle) -> Option<PersistedStats> {
    println!("Loading stats from store...");

    let store = match app.store("stats.json") {
        Ok(s) => s,
        Err(e) => {
            println!("Failed to open stats store: {}", e);
            return None;
        }
    };

    if let Some(stats_value) = store.get("lifetime_stats") {
        println!("Found lifetime stats in store: {:?}", stats_value);
        if let Ok(persisted) = serde_json::from_value::<PersistedStats>(stats_value.clone()) {
            println!("Loaded lifetime stats successfully");
            return Some(persisted);
        } else {
            println!("Failed to parse lifetime stats");
        }
    } else {
        println!("No lifetime stats found in store");
    }

    None
}

#[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_cost_calculation_opus_45() {
        let cost = calculate_cost(1000, 2000, "claude-opus-4-5-20251101");
        // Same pricing as Opus 4
        assert!((cost - 0.165).abs() < 0.0001);
    }

    #[test]
    fn test_cost_calculation_haiku() {
        let cost = calculate_cost(1000, 2000, "claude-3-5-haiku-20241022");
        // 1000 input * $1/M = $0.001
        // 2000 output * $5/M = $0.010
        // Total = $0.011
        assert!((cost - 0.011).abs() < 0.0001);
    }

    #[test]
    fn test_cost_calculation_unknown_defaults_to_sonnet() {
        let cost = calculate_cost(1000, 2000, "some-unknown-model");
        // Should default to Sonnet pricing
        assert!((cost - 0.033).abs() < 0.0001);
    }

    #[test]
    fn test_cost_calculation_legacy_sonnet() {
        let cost = calculate_cost(1000, 2000, "claude-3-5-sonnet-20241022");
        // Same as Sonnet 4 pricing
        assert!((cost - 0.033).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_usage_stats_multiple_accumulations() {
        let mut stats = UsageStats::new();
        stats.add_usage(1000, 1000, "claude-sonnet-4-20250514");
        stats.add_usage(500, 500, "claude-sonnet-4-20250514");

        assert_eq!(stats.total_input_tokens, 1500);
        assert_eq!(stats.total_output_tokens, 1500);
        assert_eq!(stats.session_input_tokens, 1500);
        assert_eq!(stats.session_output_tokens, 1500);
    }

    #[test]
    fn test_usage_stats_model_updated() {
        let mut stats = UsageStats::new();
        stats.add_usage(1000, 1000, "claude-sonnet-4-20250514");
        assert_eq!(stats.model, Some("claude-sonnet-4-20250514".to_string()));

        stats.add_usage(500, 500, "claude-opus-4-20250514");
        assert_eq!(stats.model, Some("claude-opus-4-20250514".to_string()));
    }

    #[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);
    }

    #[test]
    fn test_session_reset_clears_session_stats() {
        let mut stats = UsageStats::new();
        stats.increment_messages();
        stats.increment_messages();
        stats.increment_code_blocks();
        stats.increment_files_edited();
        stats.increment_files_created();
        stats.increment_tool_usage("Read");

        stats.reset_session();

        assert_eq!(stats.session_messages_exchanged, 0);
        assert_eq!(stats.session_code_blocks_generated, 0);
        assert_eq!(stats.session_files_edited, 0);
        assert_eq!(stats.session_files_created, 0);
        assert!(stats.session_tools_usage.is_empty());
    }

    #[test]
    fn test_increment_messages() {
        let mut stats = UsageStats::new();
        stats.increment_messages();
        stats.increment_messages();
        stats.increment_messages();

        assert_eq!(stats.messages_exchanged, 3);
        assert_eq!(stats.session_messages_exchanged, 3);
    }

    #[test]
    fn test_increment_code_blocks() {
        let mut stats = UsageStats::new();
        stats.increment_code_blocks();
        stats.increment_code_blocks();

        assert_eq!(stats.code_blocks_generated, 2);
        assert_eq!(stats.session_code_blocks_generated, 2);
    }

    #[test]
    fn test_increment_files_edited() {
        let mut stats = UsageStats::new();
        stats.increment_files_edited();

        assert_eq!(stats.files_edited, 1);
        assert_eq!(stats.session_files_edited, 1);
    }

    #[test]
    fn test_increment_files_created() {
        let mut stats = UsageStats::new();
        stats.increment_files_created();

        assert_eq!(stats.files_created, 1);
        assert_eq!(stats.session_files_created, 1);
    }

    #[test]
    fn test_increment_tool_usage() {
        let mut stats = UsageStats::new();
        stats.increment_tool_usage("Read");
        stats.increment_tool_usage("Read");
        stats.increment_tool_usage("Write");

        assert_eq!(stats.tools_usage.get("Read"), Some(&2));
        assert_eq!(stats.tools_usage.get("Write"), Some(&1));
        assert_eq!(stats.session_tools_usage.get("Read"), Some(&2));
        assert_eq!(stats.session_tools_usage.get("Write"), Some(&1));
    }

    #[test]
    fn test_session_duration_tracking() {
        let mut stats = UsageStats::new();
        stats.session_start = Some(Instant::now());

        // Verify duration is returned (u64 is always non-negative)
        let _duration = stats.get_session_duration();
    }

    #[test]
    fn test_session_duration_without_start() {
        let mut stats = UsageStats::new();
        stats.session_start = None;
        stats.session_duration_seconds = 100;

        // Should return the stored value when no start time
        let duration = stats.get_session_duration();
        assert_eq!(duration, 100);
    }

    #[test]
    fn test_is_consecutive_day_true() {
        assert!(is_consecutive_day("2024-01-15", "2024-01-16"));
        assert!(is_consecutive_day("2024-12-31", "2025-01-01"));
    }

    #[test]
    fn test_is_consecutive_day_false() {
        assert!(!is_consecutive_day("2024-01-15", "2024-01-15")); // Same day
        assert!(!is_consecutive_day("2024-01-15", "2024-01-17")); // Gap
        assert!(!is_consecutive_day("2024-01-15", "2024-01-14")); // Backwards
    }

    #[test]
    fn test_is_consecutive_day_invalid_dates() {
        assert!(!is_consecutive_day("invalid", "2024-01-01"));
        assert!(!is_consecutive_day("2024-01-01", "invalid"));
        assert!(!is_consecutive_day("invalid", "also-invalid"));
    }

    #[test]
    fn test_persisted_stats_from_usage_stats() {
        let mut stats = UsageStats::new();
        stats.total_input_tokens = 5000;
        stats.total_output_tokens = 10000;
        stats.total_cost_usd = 1.23;
        stats.messages_exchanged = 50;
        stats.sessions_started = 5;
        stats.consecutive_days = 3;

        let persisted = PersistedStats::from(&stats);

        assert_eq!(persisted.total_input_tokens, 5000);
        assert_eq!(persisted.total_output_tokens, 10000);
        assert_eq!(persisted.total_cost_usd, 1.23);
        assert_eq!(persisted.messages_exchanged, 50);
        assert_eq!(persisted.sessions_started, 5);
        assert_eq!(persisted.consecutive_days, 3);
    }

    #[test]
    fn test_apply_persisted_stats() {
        let persisted = PersistedStats {
            total_input_tokens: 10000,
            total_output_tokens: 20000,
            total_cost_usd: 5.50,
            messages_exchanged: 100,
            code_blocks_generated: 25,
            files_edited: 10,
            files_created: 5,
            tools_usage: {
                let mut map = HashMap::new();
                map.insert("Read".to_string(), 50);
                map
            },
            sessions_started: 10,
            consecutive_days: 7,
            total_days_used: 14,
            morning_sessions: 3,
            night_sessions: 2,
            last_session_date: Some("2024-06-15".to_string()),
        };

        let mut stats = UsageStats::new();
        stats.apply_persisted(persisted);

        assert_eq!(stats.total_input_tokens, 10000);
        assert_eq!(stats.total_output_tokens, 20000);
        assert_eq!(stats.total_cost_usd, 5.50);
        assert_eq!(stats.messages_exchanged, 100);
        assert_eq!(stats.tools_usage.get("Read"), Some(&50));
        assert_eq!(stats.consecutive_days, 7);
        assert_eq!(stats.morning_sessions, 3);
        assert_eq!(stats.last_session_date, Some("2024-06-15".to_string()));
    }

    #[test]
    fn test_usage_stats_default() {
        let stats = UsageStats::default();

        assert_eq!(stats.total_input_tokens, 0);
        assert_eq!(stats.total_output_tokens, 0);
        assert_eq!(stats.total_cost_usd, 0.0);
        assert!(stats.model.is_none());
    }

    #[test]
    fn test_persisted_stats_default() {
        let persisted = PersistedStats::default();

        assert_eq!(persisted.total_input_tokens, 0);
        assert!(persisted.last_session_date.is_none());
    }

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

        // UsageStats should be serializable (for events)
        let json = serde_json::to_string(&stats).expect("Failed to serialize");
        assert!(json.contains("total_input_tokens"));
        assert!(json.contains("1000"));
    }

    #[test]
    fn test_persisted_stats_serialization() {
        let persisted = PersistedStats {
            total_input_tokens: 1234,
            total_output_tokens: 5678,
            total_cost_usd: 0.99,
            ..Default::default()
        };

        let json = serde_json::to_string(&persisted).expect("Failed to serialize");
        let parsed: PersistedStats = serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(parsed.total_input_tokens, 1234);
        assert_eq!(parsed.total_output_tokens, 5678);
        assert!((parsed.total_cost_usd - 0.99).abs() < 0.0001);
    }

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

        let event = StatsUpdateEvent { stats };
        let json = serde_json::to_string(&event).expect("Failed to serialize");

        assert!(json.contains("stats"));
        assert!(json.contains("total_input_tokens"));
    }
}