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

use chrono::{Datelike, Local, NaiveDate, Weekday};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Represents a single day's cost data
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct DailyCost {
    pub date: String, // ISO date string (YYYY-MM-DD)
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cost_usd: f64,
    pub messages_sent: u64,
    pub sessions_count: u64,
}

/// Historical cost tracking data
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CostHistory {
    /// Daily costs indexed by date string (YYYY-MM-DD)
    pub daily_costs: HashMap<String, DailyCost>,
    /// Cost alert thresholds
    pub daily_alert_threshold: Option<f64>,
    pub weekly_alert_threshold: Option<f64>,
    pub monthly_alert_threshold: Option<f64>,
    /// Whether alerts have been triggered today
    pub daily_alert_triggered: bool,
    pub weekly_alert_triggered: bool,
    pub monthly_alert_triggered: bool,
    pub last_alert_reset_date: Option<String>,
}

impl CostHistory {
    pub fn new() -> Self {
        Self::default()
    }

    /// Get today's date as a string
    fn today_str() -> String {
        Local::now().format("%Y-%m-%d").to_string()
    }

    /// Get the start of the current week (Monday)
    fn week_start() -> NaiveDate {
        let today = Local::now().date_naive();
        let days_since_monday = today.weekday().num_days_from_monday();
        today - chrono::Duration::days(days_since_monday as i64)
    }

    /// Get the start of the current month
    fn month_start() -> NaiveDate {
        let today = Local::now().date_naive();
        NaiveDate::from_ymd_opt(today.year(), today.month(), 1).unwrap_or(today)
    }

    /// Add cost for today
    pub fn add_cost(&mut self, input_tokens: u64, output_tokens: u64, cost_usd: f64) {
        let today = Self::today_str();

        // Reset alert flags if it's a new day
        if self.last_alert_reset_date.as_ref() != Some(&today) {
            self.daily_alert_triggered = false;
            // Reset weekly on Monday
            if Local::now().weekday() == Weekday::Mon {
                self.weekly_alert_triggered = false;
            }
            // Reset monthly on the 1st
            if Local::now().day() == 1 {
                self.monthly_alert_triggered = false;
            }
            self.last_alert_reset_date = Some(today.clone());
        }

        let daily = self.daily_costs.entry(today).or_default();
        daily.input_tokens += input_tokens;
        daily.output_tokens += output_tokens;
        daily.cost_usd += cost_usd;
        daily.messages_sent += 1;
    }

    /// Increment session count for today
    pub fn increment_sessions(&mut self) {
        let today = Self::today_str();
        let daily = self.daily_costs.entry(today.clone()).or_insert_with(|| DailyCost {
            date: today,
            ..Default::default()
        });
        daily.sessions_count += 1;
    }

    /// Get today's cost
    pub fn get_today_cost(&self) -> f64 {
        self.daily_costs
            .get(&Self::today_str())
            .map(|d| d.cost_usd)
            .unwrap_or(0.0)
    }

    /// Get this week's cost (Monday to Sunday)
    pub fn get_week_cost(&self) -> f64 {
        let week_start = Self::week_start();
        self.daily_costs
            .values()
            .filter(|d| {
                NaiveDate::parse_from_str(&d.date, "%Y-%m-%d")
                    .map(|date| date >= week_start)
                    .unwrap_or(false)
            })
            .map(|d| d.cost_usd)
            .sum()
    }

    /// Get this month's cost
    pub fn get_month_cost(&self) -> f64 {
        let month_start = Self::month_start();
        self.daily_costs
            .values()
            .filter(|d| {
                NaiveDate::parse_from_str(&d.date, "%Y-%m-%d")
                    .map(|date| date >= month_start)
                    .unwrap_or(false)
            })
            .map(|d| d.cost_usd)
            .sum()
    }

    /// Get cost summary for a date range
    pub fn get_summary(&self, days: u32) -> CostSummary {
        let today = Local::now().date_naive();
        let start_date = today - chrono::Duration::days(days as i64 - 1);

        let mut total_input_tokens = 0u64;
        let mut total_output_tokens = 0u64;
        let mut total_cost = 0.0f64;
        let mut total_messages = 0u64;
        let mut total_sessions = 0u64;
        let mut daily_breakdown = Vec::new();

        for i in 0..days {
            let date = start_date + chrono::Duration::days(i as i64);
            let date_str = date.format("%Y-%m-%d").to_string();

            if let Some(daily) = self.daily_costs.get(&date_str) {
                total_input_tokens += daily.input_tokens;
                total_output_tokens += daily.output_tokens;
                total_cost += daily.cost_usd;
                total_messages += daily.messages_sent;
                total_sessions += daily.sessions_count;
                daily_breakdown.push(daily.clone());
            } else {
                daily_breakdown.push(DailyCost {
                    date: date_str,
                    ..Default::default()
                });
            }
        }

        CostSummary {
            period_days: days,
            total_input_tokens,
            total_output_tokens,
            total_cost,
            total_messages,
            total_sessions,
            average_daily_cost: if days > 0 { total_cost / days as f64 } else { 0.0 },
            daily_breakdown,
        }
    }

    /// Check if any alert thresholds are exceeded and return which ones
    pub fn check_alerts(&mut self) -> Vec<CostAlert> {
        let mut alerts = Vec::new();

        if let Some(threshold) = self.daily_alert_threshold {
            let today_cost = self.get_today_cost();
            if today_cost >= threshold && !self.daily_alert_triggered {
                self.daily_alert_triggered = true;
                alerts.push(CostAlert {
                    alert_type: AlertType::Daily,
                    threshold,
                    current_cost: today_cost,
                });
            }
        }

        if let Some(threshold) = self.weekly_alert_threshold {
            let week_cost = self.get_week_cost();
            if week_cost >= threshold && !self.weekly_alert_triggered {
                self.weekly_alert_triggered = true;
                alerts.push(CostAlert {
                    alert_type: AlertType::Weekly,
                    threshold,
                    current_cost: week_cost,
                });
            }
        }

        if let Some(threshold) = self.monthly_alert_threshold {
            let month_cost = self.get_month_cost();
            if month_cost >= threshold && !self.monthly_alert_triggered {
                self.monthly_alert_triggered = true;
                alerts.push(CostAlert {
                    alert_type: AlertType::Monthly,
                    threshold,
                    current_cost: month_cost,
                });
            }
        }

        alerts
    }

    /// Set alert thresholds
    pub fn set_alert_thresholds(
        &mut self,
        daily: Option<f64>,
        weekly: Option<f64>,
        monthly: Option<f64>,
    ) {
        self.daily_alert_threshold = daily;
        self.weekly_alert_threshold = weekly;
        self.monthly_alert_threshold = monthly;
    }

    /// Clean up old data (keep last N days)
    #[allow(dead_code)]
    pub fn cleanup_old_data(&mut self, keep_days: u32) {
        let cutoff = Local::now().date_naive() - chrono::Duration::days(keep_days as i64);
        self.daily_costs.retain(|date_str, _| {
            NaiveDate::parse_from_str(date_str, "%Y-%m-%d")
                .map(|date| date >= cutoff)
                .unwrap_or(false)
        });
    }

    /// Export to CSV format
    pub fn export_csv(&self, days: u32) -> String {
        let summary = self.get_summary(days);
        let mut csv = String::from("Date,Input Tokens,Output Tokens,Cost (USD),Messages,Sessions\n");

        for daily in &summary.daily_breakdown {
            csv.push_str(&format!(
                "{},{},{},{:.4},{},{}\n",
                daily.date,
                daily.input_tokens,
                daily.output_tokens,
                daily.cost_usd,
                daily.messages_sent,
                daily.sessions_count
            ));
        }

        // Add totals row
        csv.push_str(&format!(
            "TOTAL,{},{},{:.4},{},{}\n",
            summary.total_input_tokens,
            summary.total_output_tokens,
            summary.total_cost,
            summary.total_messages,
            summary.total_sessions
        ));

        csv
    }
}

/// Cost summary for a period
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostSummary {
    pub period_days: u32,
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_cost: f64,
    pub total_messages: u64,
    pub total_sessions: u64,
    pub average_daily_cost: f64,
    pub daily_breakdown: Vec<DailyCost>,
}

/// Alert types
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum AlertType {
    Daily,
    Weekly,
    Monthly,
}

/// Cost alert notification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostAlert {
    pub alert_type: AlertType,
    pub threshold: f64,
    pub current_cost: f64,
}

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

    #[test]
    fn test_add_cost() {
        let mut history = CostHistory::new();
        history.add_cost(1000, 500, 0.05);

        let today_cost = history.get_today_cost();
        assert!((today_cost - 0.05).abs() < 0.0001);
    }

    #[test]
    fn test_accumulate_daily_cost() {
        let mut history = CostHistory::new();
        history.add_cost(1000, 500, 0.05);
        history.add_cost(2000, 1000, 0.10);

        let today_cost = history.get_today_cost();
        assert!((today_cost - 0.15).abs() < 0.0001);
    }

    #[test]
    fn test_summary() {
        let mut history = CostHistory::new();
        history.add_cost(1000, 500, 0.05);

        let summary = history.get_summary(7);
        assert_eq!(summary.period_days, 7);
        assert!((summary.total_cost - 0.05).abs() < 0.0001);
    }

    #[test]
    fn test_daily_alert() {
        let mut history = CostHistory::new();
        history.set_alert_thresholds(Some(0.10), None, None);

        history.add_cost(1000, 500, 0.05);
        let alerts = history.check_alerts();
        assert!(alerts.is_empty());

        history.add_cost(1000, 500, 0.06);
        let alerts = history.check_alerts();
        assert_eq!(alerts.len(), 1);
        assert_eq!(alerts[0].alert_type, AlertType::Daily);
    }

    #[test]
    fn test_alert_only_triggers_once() {
        let mut history = CostHistory::new();
        history.set_alert_thresholds(Some(0.10), None, None);

        history.add_cost(1000, 500, 0.15);
        let alerts = history.check_alerts();
        assert_eq!(alerts.len(), 1);

        // Second check should not trigger again
        let alerts = history.check_alerts();
        assert!(alerts.is_empty());
    }

    #[test]
    fn test_export_csv() {
        let mut history = CostHistory::new();
        history.add_cost(1000, 500, 0.05);

        let csv = history.export_csv(1);
        assert!(csv.contains("Date,Input Tokens"));
        assert!(csv.contains("TOTAL"));
    }

    #[test]
    fn test_increment_sessions() {
        let mut history = CostHistory::new();
        history.increment_sessions();
        history.increment_sessions();

        let summary = history.get_summary(1);
        assert_eq!(summary.total_sessions, 2);
    }
}