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;

/// Per-tool token usage statistics
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ToolTokenStats {
    pub call_count: u64,
    pub estimated_input_tokens: u64,
    pub estimated_output_tokens: u64,
}

impl ToolTokenStats {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn increment_call(&mut self) {
        self.call_count += 1;
    }

    pub fn add_tokens(&mut self, input: u64, output: u64) {
        self.estimated_input_tokens += input;
        self.estimated_output_tokens += output;
    }

    #[allow(dead_code)]
    pub fn total_tokens(&self) -> u64 {
        self.estimated_input_tokens + self.estimated_output_tokens
    }
}

/// Warning levels for context window utilisation
#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ContextWarning {
    /// 50-74% utilisation - conversation is getting long
    Moderate,
    /// 75-89% utilisation - consider summarising
    High,
    /// 90%+ utilisation - approaching limit
    Critical,
}

/// Budget status indicating whether user is within their limits
#[allow(dead_code)]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum BudgetStatus {
    /// Within budget, no concerns
    Ok,
    /// Approaching budget limit (warning threshold reached)
    Warning {
        budget_type: BudgetType,
        percent_used: f32,
    },
    /// Budget exceeded
    Exceeded { budget_type: BudgetType },
}

/// Type of budget limit
#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum BudgetType {
    Token,
    Cost,
}

impl ContextWarning {
    #[allow(dead_code)]
    pub fn message(&self) -> &'static str {
        match self {
            ContextWarning::Moderate => "Context window is 50%+ full. Consider starting a new conversation for better performance.",
            ContextWarning::High => "Context window is 75%+ full. Responses may degrade. Consider summarising or starting fresh.",
            ContextWarning::Critical => "Context window is nearly full (90%+)! Start a new conversation to avoid errors.",
        }
    }
}

/// Get the context window limit (in tokens) for a given model
fn get_context_window_limit(model: &str) -> u64 {
    match model {
        // Claude 4.6 family - 200K standard (1M beta available via header)
        "claude-opus-4-6" => 200_000,
        // Claude 4.5 family - 200K standard context
        "claude-opus-4-5-20251101"
        | "claude-sonnet-4-5-20250929"
        | "claude-haiku-4-5-20251001" => 200_000,
        // Claude 4.x family - 200K standard context
        "claude-opus-4-1-20250805"
        | "claude-opus-4-20250514"
        | "claude-sonnet-4-20250514" => 200_000,
        // Claude 3.x family
        "claude-3-7-sonnet-20250219"
        | "claude-3-5-sonnet-20241022"
        | "claude-3-5-sonnet-20240620"
        | "claude-3-5-haiku-20241022"
        | "claude-3-opus-20240229"
        | "claude-3-sonnet-20240229"
        | "claude-3-haiku-20240307" => 200_000,
        // Default to 200K for unknown Claude models
        _ if model.starts_with("claude") => 200_000,
        // For non-Claude models (Ollama, etc.), use a conservative default
        _ => 128_000,
    }
}

#[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, ToolTokenStats>,
    pub session_tools_usage: HashMap<String, ToolTokenStats>,
    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

    // Context window tracking
    pub context_tokens_used: u64,
    pub context_window_limit: u64,
    pub context_utilisation_percent: f32,

    // Cache analytics (tracks potential savings from repeated tool calls)
    pub potential_cache_hits: u64,
    pub potential_cache_savings_tokens: u64,

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

    // Track current in-flight request for cost estimation on interrupt
    #[serde(skip)]
    pub current_request_input: Option<String>,
    #[serde(skip)]
    pub current_request_output_chars: u64,
    #[serde(skip)]
    pub current_request_thinking_chars: u64,
    #[serde(skip)]
    pub current_request_tools: Vec<String>,
}

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,
        cache_creation_tokens: Option<u64>,
        cache_read_tokens: Option<u64>,
    ) {
        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,
            cache_creation_tokens,
            cache_read_tokens,
        );
        self.total_cost_usd += cost;
        self.session_cost_usd += cost;

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

        // Update context window tracking
        self.update_context_tracking(model);
    }

    pub fn update_context_tracking(&mut self, model: &str) {
        // Get context window limit for the current model
        self.context_window_limit = get_context_window_limit(model);

        // Context tokens = input tokens (the prompt/context sent to the model)
        // We track cumulative session input as a proxy for context growth
        self.context_tokens_used = self.session_input_tokens;

        // Calculate utilisation percentage
        if self.context_window_limit > 0 {
            self.context_utilisation_percent =
                (self.context_tokens_used as f32 / self.context_window_limit as f32) * 100.0;
        }
    }

    #[allow(dead_code)]
    pub fn get_context_warning(&self) -> Option<ContextWarning> {
        if self.context_utilisation_percent >= 90.0 {
            Some(ContextWarning::Critical)
        } else if self.context_utilisation_percent >= 75.0 {
            Some(ContextWarning::High)
        } else if self.context_utilisation_percent >= 50.0 {
            Some(ContextWarning::Moderate)
        } else {
            None
        }
    }

    #[allow(dead_code)]
    pub fn estimate_remaining_tokens(&self) -> u64 {
        self.context_window_limit
            .saturating_sub(self.context_tokens_used)
    }

    /// Check budget status given current usage and budget settings
    #[allow(dead_code)]
    pub fn check_budget(
        &self,
        budget_enabled: bool,
        token_budget: Option<u64>,
        cost_budget: Option<f64>,
        warning_threshold: f32,
    ) -> BudgetStatus {
        if !budget_enabled {
            return BudgetStatus::Ok;
        }

        let session_tokens = self.session_input_tokens + self.session_output_tokens;

        // Check token budget
        if let Some(limit) = token_budget {
            if session_tokens >= limit {
                return BudgetStatus::Exceeded {
                    budget_type: BudgetType::Token,
                };
            }
            let percent_used = session_tokens as f32 / limit as f32;
            if percent_used >= warning_threshold {
                return BudgetStatus::Warning {
                    budget_type: BudgetType::Token,
                    percent_used: percent_used * 100.0,
                };
            }
        }

        // Check cost budget
        if let Some(limit) = cost_budget {
            if self.session_cost_usd >= limit {
                return BudgetStatus::Exceeded {
                    budget_type: BudgetType::Cost,
                };
            }
            let percent_used = (self.session_cost_usd / limit) as f32;
            if percent_used >= warning_threshold {
                return BudgetStatus::Warning {
                    budget_type: BudgetType::Cost,
                    percent_used: percent_used * 100.0,
                };
            }
        }

        BudgetStatus::Ok
    }

    /// Get remaining token budget (None if no budget set)
    #[allow(dead_code)]
    pub fn get_remaining_token_budget(&self, token_budget: Option<u64>) -> Option<u64> {
        token_budget.map(|limit| {
            let used = self.session_input_tokens + self.session_output_tokens;
            limit.saturating_sub(used)
        })
    }

    /// Get remaining cost budget (None if no budget set)
    #[allow(dead_code)]
    pub fn get_remaining_cost_budget(&self, cost_budget: Option<f64>) -> Option<f64> {
        cost_budget.map(|limit| {
            if limit > self.session_cost_usd {
                limit - self.session_cost_usd
            } else {
                0.0
            }
        })
    }

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

        // Reset context window tracking
        self.context_tokens_used = 0;
        self.context_utilisation_percent = 0.0;

        // Note: Cache analytics are NOT reset here - they're cumulative across sessions
        // to show total potential savings over time

        // 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_default()
            .increment_call();
        self.session_tools_usage
            .entry(tool_name.to_string())
            .or_default()
            .increment_call();
    }

    pub fn add_tool_tokens(&mut self, tool_name: &str, input_tokens: u64, output_tokens: u64) {
        self.tools_usage
            .entry(tool_name.to_string())
            .or_default()
            .add_tokens(input_tokens, output_tokens);
        self.session_tools_usage
            .entry(tool_name.to_string())
            .or_default()
            .add_tokens(input_tokens, output_tokens);
    }

    /// Record a potential cache hit (when the same tool call is made twice)
    #[allow(dead_code)]
    pub fn add_potential_cache_hit(&mut self, tokens_saved: u64) {
        self.potential_cache_hits += 1;
        self.potential_cache_savings_tokens += tokens_saved;
    }

    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(),
            context_tokens_used: self.context_tokens_used,
            context_window_limit: self.context_window_limit,
            context_utilisation_percent: self.context_utilisation_percent,
            potential_cache_hits: self.potential_cache_hits,
            potential_cache_savings_tokens: self.potential_cache_savings_tokens,
            achievements: AchievementProgress::new(), // Dummy for copy
            current_request_input: None, // Don't copy tracking fields
            current_request_output_chars: 0,
            current_request_thinking_chars: 0,
            current_request_tools: Vec::new(),
        };
        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 February 2026
// https://platform.claude.com/docs/en/about-claude/models/overview
// Cache pricing: https://platform.claude.com/docs/en/build-with-claude/prompt-caching
pub fn calculate_cost(
    input_tokens: u64,
    output_tokens: u64,
    model: &str,
    cache_creation_tokens: Option<u64>,
    cache_read_tokens: Option<u64>,
) -> f64 {
    let (input_price_per_million, output_price_per_million) = match model {
        // Current generation (Claude 4.6)
        "claude-opus-4-6" => (5.0, 25.0),

        // Previous generation (Claude 4.5)
        "claude-opus-4-5-20251101" => (5.0, 25.0),
        "claude-sonnet-4-5-20250929" => (3.0, 15.0),
        "claude-haiku-4-5-20251001" => (1.0, 5.0),

        // Previous generation (Claude 4.x)
        "claude-opus-4-1-20250805" => (15.0, 75.0),
        "claude-opus-4-20250514" => (15.0, 75.0),
        "claude-sonnet-4-20250514" => (3.0, 15.0),

        // Legacy (Claude 3.x)
        "claude-3-7-sonnet-20250219" => (3.0, 15.0),
        "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),
    };

    // Regular input/output tokens
    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;

    // Cache write tokens (cache creation) cost 1.25x the base input price
    let cache_write_cost = if let Some(cache_creation) = cache_creation_tokens {
        (cache_creation as f64 / 1_000_000.0) * input_price_per_million * 1.25
    } else {
        0.0
    };

    // Cache read tokens cost 0.1x (10%) the base input price
    let cache_read_cost = if let Some(cache_read) = cache_read_tokens {
        (cache_read as f64 / 1_000_000.0) * input_price_per_million * 0.1
    } else {
        0.0
    };

    input_cost + output_cost + cache_write_cost + cache_read_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, ToolTokenStats>,
    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);

    tracing::info!("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())?;

    tracing::info!("Stats saved successfully");
    Ok(())
}

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

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

    if let Some(stats_value) = store.get("lifetime_stats") {
        tracing::info!("Found lifetime stats in store: {:?}", stats_value);
        if let Ok(persisted) = serde_json::from_value::<PersistedStats>(stats_value.clone()) {
            tracing::info!("Loaded lifetime stats successfully");
            return Some(persisted);
        } else {
            tracing::error!("Failed to parse lifetime stats");
        }
    } else {
        tracing::info!("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", None, None);
        // 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", None, None);
        // 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", None, None);
        // Opus 4.5 pricing: $5/MTok input, $25/MTok output
        // 1000 input tokens = $0.005, 2000 output tokens = $0.05
        // Total = $0.055
        assert!((cost - 0.055).abs() < 0.0001);
    }

    #[test]
    fn test_cost_calculation_haiku() {
        let cost = calculate_cost(1000, 2000, "claude-3-5-haiku-20241022", None, None);
        // 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", None, None);
        // 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", None, None);
        // 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", None, None);

        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", None, None);
        stats.add_usage(500, 500, "claude-sonnet-4-20250514", None, None);

        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", None, None);
        assert_eq!(stats.model, Some("claude-sonnet-4-20250514".to_string()));

        stats.add_usage(500, 500, "claude-opus-4-20250514", None, None);
        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", None, None);
        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").map(|t| t.call_count), Some(2));
        assert_eq!(stats.tools_usage.get("Write").map(|t| t.call_count), Some(1));
        assert_eq!(stats.session_tools_usage.get("Read").map(|t| t.call_count), Some(2));
        assert_eq!(stats.session_tools_usage.get("Write").map(|t| t.call_count), Some(1));
    }

    #[test]
    fn test_add_tool_tokens() {
        let mut stats = UsageStats::new();
        stats.increment_tool_usage("Read");
        stats.add_tool_tokens("Read", 100, 50);
        stats.add_tool_tokens("Read", 200, 100);

        let read_stats = stats.tools_usage.get("Read").unwrap();
        assert_eq!(read_stats.call_count, 1);
        assert_eq!(read_stats.estimated_input_tokens, 300);
        assert_eq!(read_stats.estimated_output_tokens, 150);
        assert_eq!(read_stats.total_tokens(), 450);
    }

    #[test]
    fn test_tool_token_stats_default() {
        let tool_stats = ToolTokenStats::new();
        assert_eq!(tool_stats.call_count, 0);
        assert_eq!(tool_stats.estimated_input_tokens, 0);
        assert_eq!(tool_stats.estimated_output_tokens, 0);
        assert_eq!(tool_stats.total_tokens(), 0);
    }

    #[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(), ToolTokenStats {
                    call_count: 50,
                    estimated_input_tokens: 5000,
                    estimated_output_tokens: 2500,
                });
                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").map(|t| t.call_count), Some(50));
        assert_eq!(stats.tools_usage.get("Read").map(|t| t.estimated_input_tokens), Some(5000));
        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", None, None);
        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", None, None);

        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"));
    }

    // =====================
    // Context Window Tracking tests
    // =====================

    #[test]
    fn test_context_window_limit_claude_4() {
        assert_eq!(get_context_window_limit("claude-opus-4-5-20251101"), 200_000);
        assert_eq!(get_context_window_limit("claude-opus-4-20250514"), 200_000);
        assert_eq!(get_context_window_limit("claude-sonnet-4-20250514"), 200_000);
    }

    #[test]
    fn test_context_window_limit_claude_35() {
        assert_eq!(
            get_context_window_limit("claude-3-5-sonnet-20241022"),
            200_000
        );
        assert_eq!(
            get_context_window_limit("claude-3-5-sonnet-20240620"),
            200_000
        );
        assert_eq!(
            get_context_window_limit("claude-3-5-haiku-20241022"),
            200_000
        );
    }

    #[test]
    fn test_context_window_limit_unknown_claude() {
        assert_eq!(
            get_context_window_limit("claude-some-future-model"),
            200_000
        );
    }

    #[test]
    fn test_context_window_limit_non_claude() {
        assert_eq!(get_context_window_limit("gpt-4"), 128_000);
        assert_eq!(get_context_window_limit("llama-3"), 128_000);
        assert_eq!(get_context_window_limit("unknown-model"), 128_000);
    }

    #[test]
    fn test_context_tracking_update() {
        let mut stats = UsageStats::new();
        stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);

        assert_eq!(stats.context_tokens_used, 50_000);
        assert_eq!(stats.context_window_limit, 200_000);
        assert!((stats.context_utilisation_percent - 25.0).abs() < 0.1);
    }

    #[test]
    fn test_context_tracking_accumulates() {
        let mut stats = UsageStats::new();
        stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);
        stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);

        assert_eq!(stats.context_tokens_used, 100_000);
        assert!((stats.context_utilisation_percent - 50.0).abs() < 0.1);
    }

    #[test]
    fn test_context_warning_none() {
        let mut stats = UsageStats::new();
        stats.context_utilisation_percent = 40.0;
        assert!(stats.get_context_warning().is_none());
    }

    #[test]
    fn test_context_warning_moderate() {
        let mut stats = UsageStats::new();
        stats.context_utilisation_percent = 55.0;
        assert_eq!(stats.get_context_warning(), Some(ContextWarning::Moderate));
    }

    #[test]
    fn test_context_warning_high() {
        let mut stats = UsageStats::new();
        stats.context_utilisation_percent = 80.0;
        assert_eq!(stats.get_context_warning(), Some(ContextWarning::High));
    }

    #[test]
    fn test_context_warning_critical() {
        let mut stats = UsageStats::new();
        stats.context_utilisation_percent = 95.0;
        assert_eq!(stats.get_context_warning(), Some(ContextWarning::Critical));
    }

    #[test]
    fn test_estimate_remaining_tokens() {
        let mut stats = UsageStats::new();
        stats.context_tokens_used = 50_000;
        stats.context_window_limit = 200_000;

        assert_eq!(stats.estimate_remaining_tokens(), 150_000);
    }

    #[test]
    fn test_estimate_remaining_tokens_at_limit() {
        let mut stats = UsageStats::new();
        stats.context_tokens_used = 200_000;
        stats.context_window_limit = 200_000;

        assert_eq!(stats.estimate_remaining_tokens(), 0);
    }

    #[test]
    fn test_estimate_remaining_tokens_over_limit() {
        let mut stats = UsageStats::new();
        stats.context_tokens_used = 250_000;
        stats.context_window_limit = 200_000;

        assert_eq!(stats.estimate_remaining_tokens(), 0);
    }

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

        assert!(stats.context_tokens_used > 0);
        assert!(stats.context_utilisation_percent > 0.0);

        stats.reset_session();

        assert_eq!(stats.context_tokens_used, 0);
        assert_eq!(stats.context_utilisation_percent, 0.0);
    }

    #[test]
    fn test_context_warning_message() {
        assert_eq!(
            ContextWarning::Moderate.message(),
            "Context window is 50%+ full. Consider starting a new conversation for better performance."
        );
        assert_eq!(
            ContextWarning::High.message(),
            "Context window is 75%+ full. Responses may degrade. Consider summarising or starting fresh."
        );
        assert_eq!(
            ContextWarning::Critical.message(),
            "Context window is nearly full (90%+)! Start a new conversation to avoid errors."
        );
    }

    #[test]
    fn test_context_warning_serialization() {
        let warning = ContextWarning::Critical;
        let json = serde_json::to_string(&warning).expect("Failed to serialize");
        assert_eq!(json, "\"critical\"");

        let warning = ContextWarning::Moderate;
        let json = serde_json::to_string(&warning).expect("Failed to serialize");
        assert_eq!(json, "\"moderate\"");
    }

    // =====================
    // Budget Tracking tests
    // =====================

    #[test]
    fn test_budget_disabled_returns_ok() {
        let stats = UsageStats::new();
        let status = stats.check_budget(false, Some(1000), Some(1.0), 0.8);
        assert_eq!(status, BudgetStatus::Ok);
    }

    #[test]
    fn test_budget_no_limits_returns_ok() {
        let stats = UsageStats::new();
        let status = stats.check_budget(true, None, None, 0.8);
        assert_eq!(status, BudgetStatus::Ok);
    }

    #[test]
    fn test_token_budget_within_limit() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 500;
        stats.session_output_tokens = 300;

        let status = stats.check_budget(true, Some(10000), None, 0.8);
        assert_eq!(status, BudgetStatus::Ok);
    }

    #[test]
    fn test_token_budget_warning() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 4500;
        stats.session_output_tokens = 4000;

        let status = stats.check_budget(true, Some(10000), None, 0.8);
        match status {
            BudgetStatus::Warning {
                budget_type,
                percent_used,
            } => {
                assert_eq!(budget_type, BudgetType::Token);
                assert!(percent_used >= 80.0);
            }
            _ => panic!("Expected Warning status"),
        }
    }

    #[test]
    fn test_token_budget_exceeded() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 6000;
        stats.session_output_tokens = 5000;

        let status = stats.check_budget(true, Some(10000), None, 0.8);
        assert_eq!(
            status,
            BudgetStatus::Exceeded {
                budget_type: BudgetType::Token
            }
        );
    }

    #[test]
    fn test_cost_budget_within_limit() {
        let mut stats = UsageStats::new();
        stats.session_cost_usd = 0.50;

        let status = stats.check_budget(true, None, Some(5.0), 0.8);
        assert_eq!(status, BudgetStatus::Ok);
    }

    #[test]
    fn test_cost_budget_warning() {
        let mut stats = UsageStats::new();
        stats.session_cost_usd = 4.25;

        let status = stats.check_budget(true, None, Some(5.0), 0.8);
        match status {
            BudgetStatus::Warning {
                budget_type,
                percent_used,
            } => {
                assert_eq!(budget_type, BudgetType::Cost);
                assert!(percent_used >= 80.0);
            }
            _ => panic!("Expected Warning status"),
        }
    }

    #[test]
    fn test_cost_budget_exceeded() {
        let mut stats = UsageStats::new();
        stats.session_cost_usd = 5.50;

        let status = stats.check_budget(true, None, Some(5.0), 0.8);
        assert_eq!(
            status,
            BudgetStatus::Exceeded {
                budget_type: BudgetType::Cost
            }
        );
    }

    #[test]
    fn test_token_budget_takes_priority() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 12000;
        stats.session_output_tokens = 0;
        stats.session_cost_usd = 0.01;

        // Token budget exceeded, cost budget OK
        let status = stats.check_budget(true, Some(10000), Some(5.0), 0.8);
        assert_eq!(
            status,
            BudgetStatus::Exceeded {
                budget_type: BudgetType::Token
            }
        );
    }

    #[test]
    fn test_remaining_token_budget() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 3000;
        stats.session_output_tokens = 2000;

        assert_eq!(stats.get_remaining_token_budget(Some(10000)), Some(5000));
        assert_eq!(stats.get_remaining_token_budget(None), None);
    }

    #[test]
    fn test_remaining_token_budget_exceeded() {
        let mut stats = UsageStats::new();
        stats.session_input_tokens = 8000;
        stats.session_output_tokens = 5000;

        assert_eq!(stats.get_remaining_token_budget(Some(10000)), Some(0));
    }

    #[test]
    fn test_remaining_cost_budget() {
        let mut stats = UsageStats::new();
        stats.session_cost_usd = 2.50;

        let remaining = stats.get_remaining_cost_budget(Some(5.0));
        assert!(remaining.is_some());
        assert!((remaining.unwrap() - 2.50).abs() < 0.001);
        assert_eq!(stats.get_remaining_cost_budget(None), None);
    }

    #[test]
    fn test_remaining_cost_budget_exceeded() {
        let mut stats = UsageStats::new();
        stats.session_cost_usd = 6.0;

        let remaining = stats.get_remaining_cost_budget(Some(5.0));
        assert!(remaining.is_some());
        assert!((remaining.unwrap() - 0.0).abs() < 0.001);
    }

    #[test]
    fn test_budget_status_serialization() {
        let status = BudgetStatus::Warning {
            budget_type: BudgetType::Token,
            percent_used: 85.5,
        };
        let json = serde_json::to_string(&status).expect("Failed to serialize");
        assert!(json.contains("warning"));
        assert!(json.contains("token"));

        let status = BudgetStatus::Exceeded {
            budget_type: BudgetType::Cost,
        };
        let json = serde_json::to_string(&status).expect("Failed to serialize");
        assert!(json.contains("exceeded"));
        assert!(json.contains("cost"));
    }

    #[test]
    fn test_budget_type_serialization() {
        let token = BudgetType::Token;
        let json = serde_json::to_string(&token).expect("Failed to serialize");
        assert_eq!(json, "\"token\"");

        let cost = BudgetType::Cost;
        let json = serde_json::to_string(&cost).expect("Failed to serialize");
        assert_eq!(json, "\"cost\"");
    }
}