tokmesh-cli 0.1.1

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::path::PathBuf;

use anyhow::Result;
use chrono::{Datelike, NaiveDate, NaiveDateTime, Timelike};
use tokio::runtime::{Handle, Runtime};

use tokmesh_core::sessions::UnifiedMessage;
use tokmesh_core::{
    model_name_for_grouping, normalize_model_for_grouping, parse_local_unified_messages, sessions,
    ClientId, GroupBy, LocalParseOptions, ModelPerformance,
};

/// Returns the scanner settings that `DataLoader` should use when building
/// `LocalParseOptions`. Under `#[cfg(test)]` this intentionally ignores
/// `~/.config/tokmesh/settings.json` so data-loader unit tests stay
/// hermetic across developer machines; production builds still honor
/// user-configured paths.
#[cfg(not(test))]
fn data_loader_scanner_settings() -> tokmesh_core::scanner::ScannerSettings {
    crate::tui::settings::load_scanner_settings()
}

#[cfg(test)]
fn data_loader_scanner_settings() -> tokmesh_core::scanner::ScannerSettings {
    tokmesh_core::scanner::ScannerSettings::default()
}

#[derive(Debug, Clone, Default)]
pub struct TokenBreakdown {
    pub input: u64,
    pub output: u64,
    pub cache_read: u64,
    pub cache_write: u64,
    pub reasoning: u64,
}

impl TokenBreakdown {
    pub fn total(&self) -> u64 {
        self.input
            .saturating_add(self.output)
            .saturating_add(self.cache_read)
            .saturating_add(self.cache_write)
            .saturating_add(self.reasoning)
    }
}

#[derive(Debug, Clone)]
pub struct ModelUsage {
    pub model: String,
    pub color_key: String,
    pub provider: String,
    pub client: String,
    pub workspace_key: Option<String>,
    pub workspace_label: Option<String>,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub performance: ModelPerformance,
    pub session_count: u32,
}

#[derive(Debug, Clone)]
pub struct AgentUsage {
    pub agent: String,
    pub clients: String,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub message_count: u32,
}

#[derive(Debug, Clone)]
pub struct DailyModelInfo {
    /// API provider identifier (e.g. "anthropic", "openai").
    ///
    /// **Caveat**: For `GroupBy::Model`, `GroupBy::ClientModel`, and
    /// `GroupBy::WorkspaceModel`, multiple providers may be merged into a
    /// single daily model entry.  In that case this field retains whichever
    /// provider was seen first and is **not** authoritative.  Only treat it
    /// as exact when `group_by == GroupBy::ClientProviderModel`.
    pub provider: String,
    pub display_name: String,
    pub color_key: String,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub messages: u64,
}

#[derive(Debug, Clone)]
pub struct DailySourceInfo {
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub models: BTreeMap<String, DailyModelInfo>,
}

#[derive(Debug, Clone)]
pub struct DailyUsage {
    pub date: NaiveDate,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub source_breakdown: BTreeMap<String, DailySourceInfo>,
    pub message_count: u32,
    pub turn_count: u32,
}

#[derive(Debug, Clone)]
pub struct HourlyModelInfo {
    pub provider: String,
    pub display_name: String,
    pub color_key: String,
    pub tokens: TokenBreakdown,
    pub cost: f64,
}

#[derive(Debug, Clone)]
pub struct HourlyUsage {
    pub datetime: NaiveDateTime,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub clients: BTreeSet<String>,
    pub models: BTreeMap<String, HourlyModelInfo>,
    pub message_count: u32,
    pub turn_count: u32,
}

#[derive(Debug, Clone)]
pub struct MinutelyUsage {
    pub datetime: NaiveDateTime,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub clients: BTreeSet<String>,
    pub models: BTreeMap<String, HourlyModelInfo>,
    pub message_count: u32,
    pub turn_count: u32,
}

#[derive(Debug, Clone)]
pub struct MonthlyUsage {
    pub month: String,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub message_count: u32,
    pub turn_count: u32,
}

/// Per-session usage rollup. One row per `(client, session_id)` pair so the
/// Sessions tab can show cost/tokens broken down by individual conversation
/// alongside how long the session was active and when it last saw activity.
#[derive(Debug, Clone)]
pub struct SessionUsage {
    pub session_id: String,
    pub client: String,
    /// Human-readable session title when the source client stores one
    /// (e.g. OpenCode's `session.title` column). `None` for clients that
    /// don't record a title.
    pub title: Option<String>,
    /// Distinct models used across messages in this session, in first-seen
    /// order. Most sessions use a single model; a few switch mid-conversation.
    pub models: Vec<SessionModel>,
    pub tokens: TokenBreakdown,
    pub cost: f64,
    pub message_count: u32,
    pub turn_count: u32,
    /// Unix-ms timestamp of the first message observed in this session.
    /// `0` when every message lacked a usable timestamp.
    pub first_active_ms: i64,
    /// Unix-ms timestamp of the most recent message observed in this session.
    /// `0` when every message lacked a usable timestamp.
    pub last_active_ms: i64,
}

/// A model entry within a session, retaining the provider and color_key
/// needed for correct family-shade color lookup alongside the display name.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionModel {
    pub display_name: String,
    pub provider: String,
    pub color_key: String,
}

#[derive(Debug, Clone)]
pub struct ContributionDay {
    pub date: NaiveDate,
    pub tokens: u64,
    pub cost: f64,
    pub intensity: f64,
}

#[derive(Debug, Clone)]
pub struct GraphData {
    pub weeks: Vec<Vec<Option<ContributionDay>>>,
}

#[derive(Debug, Clone, Default)]
pub struct UsageData {
    pub models: Vec<ModelUsage>,
    pub agents: Vec<AgentUsage>,
    pub daily: Vec<DailyUsage>,
    pub hourly: Vec<HourlyUsage>,
    pub minutely: Vec<MinutelyUsage>,
    pub monthly: Vec<MonthlyUsage>,
    pub sessions: Vec<SessionUsage>,
    pub graph: Option<GraphData>,
    pub total_tokens: u64,
    pub total_cost: f64,
    pub loading: bool,
    pub error: Option<String>,
    pub current_streak: u32,
    pub longest_streak: u32,
}

pub struct DataLoader {
    _sessions_path: Option<PathBuf>,
    pub since: Option<String>,
    pub until: Option<String>,
    pub year: Option<String>,
    pub minutely_enabled: bool,
}

const UNKNOWN_WORKSPACE_LABEL: &str = "Unknown workspace";
const UNKNOWN_WORKSPACE_GROUP_KEY: &str = "\0unknown-workspace";

fn workspace_bucket(msg: &UnifiedMessage) -> (String, Option<String>, String) {
    match (&msg.workspace_key, &msg.workspace_label) {
        (Some(key), Some(label)) => (key.clone(), Some(key.clone()), label.clone()),
        (Some(key), None) => (
            key.clone(),
            Some(key.clone()),
            tokmesh_core::sessions::workspace_label_from_key(key)
                .unwrap_or_else(|| UNKNOWN_WORKSPACE_LABEL.to_string()),
        ),
        _ => (
            UNKNOWN_WORKSPACE_GROUP_KEY.to_string(),
            None,
            UNKNOWN_WORKSPACE_LABEL.to_string(),
        ),
    }
}

fn positive_unified_token_total(tokens: &tokmesh_core::TokenBreakdown) -> i64 {
    // saturating_add (mirrors tokmesh_core::TokenBreakdown::total) so a
    // clamped (i64::MAX) bucket from a corrupt source can't overflow the
    // per-message sum.
    tokens
        .input
        .max(0)
        .saturating_add(tokens.output.max(0))
        .saturating_add(tokens.cache_read.max(0))
        .saturating_add(tokens.cache_write.max(0))
        .saturating_add(tokens.reasoning.max(0))
}

fn workspace_model_display_label(workspace_label: &str, model: &str) -> String {
    format!("{workspace_label} / {model}")
}

fn workspace_model_daily_key(workspace_group_key: &str, model: &str) -> String {
    format!(
        "{}:{workspace_group_key}:{model}",
        workspace_group_key.len()
    )
}

fn daily_source_model_key(
    group_by: &GroupBy,
    workspace_group_key: &str,
    provider_id: &str,
    model: &str,
) -> String {
    match group_by {
        GroupBy::WorkspaceModel => workspace_model_daily_key(workspace_group_key, model),
        GroupBy::ClientProviderModel => format!("{provider_id}:{model}"),
        GroupBy::Model | GroupBy::ClientModel | GroupBy::Session | GroupBy::ClientSession => {
            model.to_string()
        }
    }
}

fn daily_source_model_display_name(
    group_by: &GroupBy,
    workspace_label: &str,
    provider_id: &str,
    model: &str,
) -> String {
    match group_by {
        GroupBy::WorkspaceModel => workspace_model_display_label(workspace_label, model),
        GroupBy::ClientProviderModel => format!("{provider_id} / {model}"),
        GroupBy::Model | GroupBy::ClientModel | GroupBy::Session | GroupBy::ClientSession => {
            model.to_string()
        }
    }
}

fn model_color_key(group_by: &GroupBy, _provider_id: &str, model: &str) -> String {
    match group_by {
        GroupBy::ClientProviderModel => model.to_string(),
        GroupBy::Model
        | GroupBy::ClientModel
        | GroupBy::WorkspaceModel
        | GroupBy::Session
        | GroupBy::ClientSession => model.to_string(),
    }
}

fn hourly_model_key(group_by: &GroupBy, provider_id: &str, model: &str) -> String {
    match group_by {
        GroupBy::ClientProviderModel => format!("{provider_id}:{model}"),
        GroupBy::Model
        | GroupBy::ClientModel
        | GroupBy::WorkspaceModel
        | GroupBy::Session
        | GroupBy::ClientSession => model.to_string(),
    }
}

fn hourly_model_display_name(group_by: &GroupBy, provider_id: &str, model: &str) -> String {
    match group_by {
        GroupBy::ClientProviderModel => format!("{provider_id} / {model}"),
        GroupBy::Model
        | GroupBy::ClientModel
        | GroupBy::WorkspaceModel
        | GroupBy::Session
        | GroupBy::ClientSession => model.to_string(),
    }
}

impl DataLoader {
    pub fn new(sessions_path: Option<PathBuf>) -> Self {
        Self {
            _sessions_path: sessions_path,
            since: None,
            until: None,
            year: None,
            minutely_enabled: false,
        }
    }

    pub fn with_filters(
        sessions_path: Option<PathBuf>,
        since: Option<String>,
        until: Option<String>,
        year: Option<String>,
    ) -> Self {
        Self {
            _sessions_path: sessions_path,
            since,
            until,
            year,
            minutely_enabled: false,
        }
    }

    pub fn with_minutely_enabled(mut self, enabled: bool) -> Self {
        self.minutely_enabled = enabled;
        self
    }

    pub fn load(
        &self,
        enabled_clients: &[ClientId],
        group_by: &GroupBy,
        include_synthetic: bool,
    ) -> Result<UsageData> {
        let home = dirs::home_dir()
            .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))?
            .to_string_lossy()
            .to_string();

        let mut sources: Vec<String> = enabled_clients
            .iter()
            .map(|client| client.as_str().to_string())
            .collect();
        if include_synthetic {
            sources.push("synthetic".to_string());
        }

        let opts = LocalParseOptions {
            home_dir: Some(home),
            use_env_roots: true,
            clients: Some(sources),
            since: self.since.clone(),
            until: self.until.clone(),
            year: self.year.clone(),
            scanner_settings: data_loader_scanner_settings(),
        };

        let messages = if Handle::try_current().is_ok() {
            std::thread::scope(|s| {
                s.spawn(|| {
                    let rt = Runtime::new().map_err(|e| e.to_string())?;
                    rt.block_on(parse_local_unified_messages(opts))
                })
                .join()
                .unwrap_or_else(|_| Err("data loader thread panicked".to_string()))
            })
        } else {
            Runtime::new()?.block_on(parse_local_unified_messages(opts))
        }
        .map_err(anyhow::Error::msg)?;

        self.aggregate_messages(messages, group_by)
    }

    #[cfg(test)]
    #[allow(dead_code)]
    fn load_with_pricing(
        &self,
        enabled_clients: &[ClientId],
        group_by: &GroupBy,
        include_synthetic: bool,
        pricing: &tokmesh_core::pricing::PricingService,
    ) -> Result<UsageData> {
        let home = dirs::home_dir()
            .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))?
            .to_string_lossy()
            .to_string();

        let mut sources: Vec<String> = enabled_clients
            .iter()
            .map(|client| client.as_str().to_string())
            .collect();
        if include_synthetic {
            sources.push("synthetic".to_string());
        }

        let opts = LocalParseOptions {
            home_dir: Some(home),
            clients: Some(sources),
            since: self.since.clone(),
            until: self.until.clone(),
            year: self.year.clone(),
            use_env_roots: false,
            scanner_settings: data_loader_scanner_settings(),
        };

        let messages = if Handle::try_current().is_ok() {
            std::thread::scope(|s| {
                s.spawn(|| {
                    let rt = Runtime::new().map_err(|e| e.to_string())?;
                    rt.block_on(tokmesh_core::parse_local_unified_messages_with_pricing(
                        opts,
                        Some(pricing),
                    ))
                })
                .join()
                .unwrap_or_else(|_| Err("data loader thread panicked".to_string()))
            })
        } else {
            Runtime::new()?.block_on(tokmesh_core::parse_local_unified_messages_with_pricing(
                opts,
                Some(pricing),
            ))
        }
        .map_err(anyhow::Error::msg)?;

        self.aggregate_messages(messages, group_by)
    }

    fn aggregate_messages(
        &self,
        messages: Vec<UnifiedMessage>,
        group_by: &GroupBy,
    ) -> Result<UsageData> {
        let mut model_map: HashMap<String, ModelUsage> = HashMap::new();
        let mut agent_map: HashMap<String, AgentUsage> = HashMap::new();
        let mut agent_clients: HashMap<String, BTreeSet<String>> = HashMap::new();
        let mut daily_map: HashMap<NaiveDate, DailyUsage> = HashMap::new();
        let mut hourly_map: HashMap<NaiveDateTime, HourlyUsage> = HashMap::new();
        let mut minutely_map: HashMap<NaiveDateTime, MinutelyUsage> = HashMap::new();
        let mut model_session_ids: HashMap<String, HashSet<String>> = HashMap::new();
        let mut session_map: HashMap<String, SessionUsage> = HashMap::new();

        for msg in &messages {
            let normalized_model =
                model_name_for_grouping(&msg.client, &msg.provider_id, &msg.model_id);
            let model_key = normalize_model_for_grouping(&msg.model_id);
            let (workspace_group_key, workspace_key, workspace_label) = workspace_bucket(msg);
            let key = match group_by {
                GroupBy::Model => normalized_model.clone(),
                GroupBy::ClientModel => format!("{}:{}", msg.client, normalized_model),
                GroupBy::ClientProviderModel => {
                    format!("{}:{}:{}", msg.client, msg.provider_id, normalized_model)
                }
                GroupBy::WorkspaceModel => {
                    format!("{}:{}", workspace_group_key, normalized_model)
                }
                GroupBy::Session => format!("{}:{}", msg.session_id, normalized_model),
                GroupBy::ClientSession => {
                    format!("{}:{}:{}", msg.client, msg.session_id, normalized_model)
                }
            };
            let merge_clients = matches!(group_by, GroupBy::Model | GroupBy::WorkspaceModel);

            let model_entry = model_map.entry(key.clone()).or_insert_with(|| ModelUsage {
                model: normalized_model.clone(),
                color_key: model_key.clone(),
                provider: msg.provider_id.clone(),
                client: msg.client.clone(),
                workspace_key: if *group_by == GroupBy::WorkspaceModel {
                    workspace_key.clone()
                } else {
                    None
                },
                workspace_label: if *group_by == GroupBy::WorkspaceModel {
                    Some(workspace_label.clone())
                } else {
                    None
                },
                tokens: TokenBreakdown::default(),
                cost: 0.0,
                performance: ModelPerformance::default(),
                session_count: 0,
            });

            if merge_clients && !model_entry.client.split(", ").any(|s| s == msg.client) {
                model_entry.client = format!("{}, {}", model_entry.client, msg.client);
            }

            if *group_by != GroupBy::ClientProviderModel
                && !model_entry
                    .provider
                    .split(", ")
                    .any(|p| p == msg.provider_id)
            {
                model_entry.provider = format!("{}, {}", model_entry.provider, msg.provider_id);
            }

            model_entry.tokens.input = model_entry
                .tokens
                .input
                .saturating_add(msg.tokens.input.max(0) as u64);
            model_entry.tokens.output = model_entry
                .tokens
                .output
                .saturating_add(msg.tokens.output.max(0) as u64);
            model_entry.tokens.cache_read = model_entry
                .tokens
                .cache_read
                .saturating_add(msg.tokens.cache_read.max(0) as u64);
            model_entry.tokens.cache_write = model_entry
                .tokens
                .cache_write
                .saturating_add(msg.tokens.cache_write.max(0) as u64);
            model_entry.tokens.reasoning = model_entry
                .tokens
                .reasoning
                .saturating_add(msg.tokens.reasoning.max(0) as u64);
            let msg_cost = if msg.cost.is_finite() && msg.cost >= 0.0 {
                msg.cost
            } else {
                0.0
            };
            model_entry.cost += msg_cost;
            model_entry
                .performance
                .record_message(positive_unified_token_total(&msg.tokens), msg.duration_ms);

            let session_key = format!("{}:{}", msg.client, msg.session_id);
            let model_sessions = model_session_ids.entry(key).or_default();
            if model_sessions.insert(session_key) {
                model_entry.session_count += 1;
            }

            if let Some(agent) = msg.agent.as_ref() {
                let normalized_agent = if msg.client == "opencode" {
                    sessions::normalize_opencode_agent_name(agent)
                } else if msg.client == "copilot" {
                    sessions::normalize_copilot_agent_name(agent)
                } else {
                    sessions::normalize_agent_name(agent)
                };
                let agent_entry = agent_map
                    .entry(normalized_agent.clone())
                    .or_insert_with(|| AgentUsage {
                        agent: normalized_agent.clone(),
                        clients: String::new(),
                        tokens: TokenBreakdown::default(),
                        cost: 0.0,
                        message_count: 0,
                    });

                agent_entry.tokens.input = agent_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                agent_entry.tokens.output = agent_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                agent_entry.tokens.cache_read = agent_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                agent_entry.tokens.cache_write = agent_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                agent_entry.tokens.reasoning = agent_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                agent_entry.cost += msg_cost;
                agent_entry.message_count = agent_entry
                    .message_count
                    .saturating_add(msg.message_count.max(0) as u32);

                agent_clients
                    .entry(normalized_agent)
                    .or_default()
                    .insert(msg.client.clone());
            }

            if let Some(date) = parse_date(&msg.date) {
                let daily_entry = daily_map.entry(date).or_insert_with(|| DailyUsage {
                    date,
                    tokens: TokenBreakdown::default(),
                    cost: 0.0,
                    source_breakdown: BTreeMap::new(),
                    message_count: 0,
                    turn_count: 0,
                });

                daily_entry.tokens.input = daily_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                daily_entry.tokens.output = daily_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                daily_entry.tokens.cache_read = daily_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                daily_entry.tokens.cache_write = daily_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                daily_entry.tokens.reasoning = daily_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                let msg_cost = if msg.cost.is_finite() && msg.cost >= 0.0 {
                    msg.cost
                } else {
                    0.0
                };
                daily_entry.cost += msg_cost;
                daily_entry.message_count += msg.message_count.max(0) as u32;
                if msg.is_turn_start {
                    daily_entry.turn_count += 1;
                }

                let source_entry = daily_entry
                    .source_breakdown
                    .entry(msg.client.clone())
                    .or_insert_with(|| DailySourceInfo {
                        tokens: TokenBreakdown::default(),
                        cost: 0.0,
                        models: BTreeMap::new(),
                    });

                source_entry.tokens.input = source_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                source_entry.tokens.output = source_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                source_entry.tokens.cache_read = source_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                source_entry.tokens.cache_write = source_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                source_entry.tokens.reasoning = source_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                source_entry.cost += msg_cost;

                let daily_model_key = daily_source_model_key(
                    group_by,
                    &workspace_group_key,
                    &msg.provider_id,
                    &normalized_model,
                );

                let model_info = source_entry
                    .models
                    .entry(daily_model_key)
                    .or_insert_with(|| DailyModelInfo {
                        provider: msg.provider_id.clone(),
                        display_name: daily_source_model_display_name(
                            group_by,
                            &workspace_label,
                            &msg.provider_id,
                            &normalized_model,
                        ),
                        color_key: model_color_key(group_by, &msg.provider_id, &model_key),
                        tokens: TokenBreakdown::default(),
                        cost: 0.0,
                        messages: 0,
                    });

                model_info.tokens.input = model_info
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                model_info.tokens.output = model_info
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                model_info.tokens.cache_read = model_info
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                model_info.tokens.cache_write = model_info
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                model_info.tokens.reasoning = model_info
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                model_info.cost += msg_cost;
                model_info.messages = model_info
                    .messages
                    .saturating_add(msg.message_count.max(0) as u64);
            }

            // Hourly aggregation: derive hour from timestamp (Unix ms),
            // falling back to msg.date 00:00 when timestamp is missing/zero
            // so we don't silently drop messages (matches CLI bucketing).
            if let Some(hour_dt) = hour_bucket_with_fallback(msg.timestamp, &msg.date) {
                let hourly_entry = hourly_map.entry(hour_dt).or_insert_with(|| HourlyUsage {
                    datetime: hour_dt,
                    tokens: TokenBreakdown::default(),
                    cost: 0.0,
                    clients: BTreeSet::new(),
                    models: BTreeMap::new(),
                    message_count: 0,
                    turn_count: 0,
                });

                hourly_entry.tokens.input = hourly_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                hourly_entry.tokens.output = hourly_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                hourly_entry.tokens.cache_read = hourly_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                hourly_entry.tokens.cache_write = hourly_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                hourly_entry.tokens.reasoning = hourly_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                let h_cost = if msg.cost.is_finite() && msg.cost >= 0.0 {
                    msg.cost
                } else {
                    0.0
                };
                hourly_entry.cost += h_cost;
                hourly_entry.message_count += msg.message_count.max(0) as u32;
                if msg.is_turn_start {
                    hourly_entry.turn_count += 1;
                }
                hourly_entry.clients.insert(msg.client.clone());

                let hourly_model_key =
                    hourly_model_key(group_by, &msg.provider_id, &normalized_model);
                let h_model = hourly_entry
                    .models
                    .entry(hourly_model_key)
                    .or_insert_with(|| HourlyModelInfo {
                        provider: msg.provider_id.clone(),
                        display_name: hourly_model_display_name(
                            group_by,
                            &msg.provider_id,
                            &normalized_model,
                        ),
                        color_key: model_color_key(group_by, &msg.provider_id, &model_key),
                        tokens: TokenBreakdown::default(),
                        cost: 0.0,
                    });
                h_model.tokens.input = h_model
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                h_model.tokens.output = h_model
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                h_model.tokens.cache_read = h_model
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                h_model.tokens.cache_write = h_model
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                h_model.tokens.reasoning = h_model
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                h_model.cost += h_cost;
            }

            // Minute aggregation: same fallback semantics as hourly, but
            // truncated to the minute. Used by the Minutely tab. Gated
            // on `Settings::minutely_tab_enabled` so users who never open
            // the tab do not pay the per-minute bucketing cost.
            let minute_bucket = if self.minutely_enabled {
                minute_bucket_with_fallback(msg.timestamp, &msg.date)
            } else {
                None
            };
            if let Some(minute_dt) = minute_bucket {
                let minutely_entry =
                    minutely_map
                        .entry(minute_dt)
                        .or_insert_with(|| MinutelyUsage {
                            datetime: minute_dt,
                            tokens: TokenBreakdown::default(),
                            cost: 0.0,
                            clients: BTreeSet::new(),
                            models: BTreeMap::new(),
                            message_count: 0,
                            turn_count: 0,
                        });

                minutely_entry.tokens.input = minutely_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                minutely_entry.tokens.output = minutely_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                minutely_entry.tokens.cache_read = minutely_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                minutely_entry.tokens.cache_write = minutely_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                minutely_entry.tokens.reasoning = minutely_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                let m_cost = if msg.cost.is_finite() && msg.cost >= 0.0 {
                    msg.cost
                } else {
                    0.0
                };
                minutely_entry.cost += m_cost;
                minutely_entry.message_count += msg.message_count.max(0) as u32;
                if msg.is_turn_start {
                    minutely_entry.turn_count += 1;
                }
                minutely_entry.clients.insert(msg.client.clone());

                let m_model_key = hourly_model_key(group_by, &msg.provider_id, &normalized_model);
                let m_model =
                    minutely_entry
                        .models
                        .entry(m_model_key)
                        .or_insert_with(|| HourlyModelInfo {
                            provider: msg.provider_id.clone(),
                            display_name: hourly_model_display_name(
                                group_by,
                                &msg.provider_id,
                                &normalized_model,
                            ),
                            color_key: model_color_key(group_by, &msg.provider_id, &model_key),
                            tokens: TokenBreakdown::default(),
                            cost: 0.0,
                        });
                m_model.tokens.input = m_model
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                m_model.tokens.output = m_model
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                m_model.tokens.cache_read = m_model
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                m_model.tokens.cache_write = m_model
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                m_model.tokens.reasoning = m_model
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                m_model.cost += m_cost;
            }

            // Session aggregation: one bucket per (client, session_id) so the
            // Sessions tab can show cost/tokens per individual conversation.
            // Skips messages with an empty session_id (some legacy/scanner
            // records lack one) rather than lumping them into a single bogus
            // "no-session" row.
            if !msg.session_id.is_empty() {
                let session_key = format!("{}:{}", msg.client, msg.session_id);
                let session_entry =
                    session_map
                        .entry(session_key)
                        .or_insert_with(|| SessionUsage {
                            session_id: msg.session_id.clone(),
                            client: msg.client.clone(),
                            title: None,
                            models: Vec::new(),
                            tokens: TokenBreakdown::default(),
                            cost: 0.0,
                            message_count: 0,
                            turn_count: 0,
                            first_active_ms: 0,
                            last_active_ms: 0,
                        });

                session_entry.tokens.input = session_entry
                    .tokens
                    .input
                    .saturating_add(msg.tokens.input.max(0) as u64);
                session_entry.tokens.output = session_entry
                    .tokens
                    .output
                    .saturating_add(msg.tokens.output.max(0) as u64);
                session_entry.tokens.cache_read = session_entry
                    .tokens
                    .cache_read
                    .saturating_add(msg.tokens.cache_read.max(0) as u64);
                session_entry.tokens.cache_write = session_entry
                    .tokens
                    .cache_write
                    .saturating_add(msg.tokens.cache_write.max(0) as u64);
                session_entry.tokens.reasoning = session_entry
                    .tokens
                    .reasoning
                    .saturating_add(msg.tokens.reasoning.max(0) as u64);
                session_entry.cost += msg_cost;
                session_entry.message_count = session_entry
                    .message_count
                    .saturating_add(msg.message_count.max(0) as u32);
                if msg.is_turn_start {
                    session_entry.turn_count += 1;
                }

                let ts = message_timestamp_ms(msg);
                if ts > 0 {
                    if session_entry.first_active_ms == 0 || ts < session_entry.first_active_ms {
                        session_entry.first_active_ms = ts;
                    }
                    if ts > session_entry.last_active_ms {
                        session_entry.last_active_ms = ts;
                    }
                }

                // Adopt the first non-empty session_title seen across the
                // session's messages. Parsers that don't populate the field
                // leave it `None` and the Sessions tab falls back to the ID.
                if session_entry.title.is_none() {
                    if let Some(ref title) = msg.session_title {
                        let trimmed = title.trim();
                        if !trimmed.is_empty() {
                            session_entry.title = Some(trimmed.to_string());
                        }
                    }
                }

                // Track distinct models in first-seen order, retaining
                // provider and color_key for correct shade-map lookups.
                if !session_entry
                    .models
                    .iter()
                    .any(|m| m.display_name == normalized_model)
                {
                    session_entry.models.push(SessionModel {
                        display_name: normalized_model.clone(),
                        provider: msg.provider_id.clone(),
                        color_key: model_key.clone(),
                    });
                }
            }
        }

        let mut models: Vec<ModelUsage> = model_map
            .into_values()
            .map(|mut model| {
                model.performance.finalize(model.tokens.total() as i64);
                model
            })
            .collect();
        models.sort_by(|a, b| {
            b.cost
                .total_cmp(&a.cost)
                .then_with(|| a.model.cmp(&b.model))
                .then_with(|| a.provider.cmp(&b.provider))
        });

        for (agent, clients) in agent_clients {
            if let Some(agent_entry) = agent_map.get_mut(&agent) {
                agent_entry.clients = clients.into_iter().collect::<Vec<_>>().join(", ");
            }
        }

        let mut agents: Vec<AgentUsage> = agent_map.into_values().collect();
        agents.sort_by(|a, b| {
            b.cost
                .total_cmp(&a.cost)
                .then_with(|| b.tokens.total().cmp(&a.tokens.total()))
                .then_with(|| a.agent.cmp(&b.agent))
        });

        let mut daily: Vec<DailyUsage> = daily_map.into_values().collect();
        daily.sort_by_key(|b| std::cmp::Reverse(b.date));

        let mut hourly: Vec<HourlyUsage> = hourly_map.into_values().collect();
        hourly.sort_by_key(|b| std::cmp::Reverse(b.datetime));

        let mut minutely: Vec<MinutelyUsage> = minutely_map.into_values().collect();
        minutely.sort_by_key(|b| std::cmp::Reverse(b.datetime));

        let monthly = aggregate_monthly_from_daily(&daily);

        let mut sessions: Vec<SessionUsage> = session_map.into_values().collect();
        sessions.sort_by(|a, b| {
            b.cost
                .total_cmp(&a.cost)
                .then_with(|| b.last_active_ms.cmp(&a.last_active_ms))
                .then_with(|| a.client.cmp(&b.client))
                .then_with(|| a.session_id.cmp(&b.session_id))
        });

        // Plain `.sum()` panics (debug) / wraps (release) on overflow across
        // many models; a single corrupt/huge bucket must not poison the
        // whole total, so fold with saturating_add like `TokenBreakdown::total`.
        let total_tokens: u64 = models
            .iter()
            .map(|m| m.tokens.total())
            .fold(0u64, u64::saturating_add);
        let total_cost: f64 = models
            .iter()
            .map(|m| if m.cost.is_finite() { m.cost } else { 0.0 })
            .sum();

        let graph = build_contribution_graph(&daily);
        let (current_streak, longest_streak) = calculate_streaks(&daily);

        Ok(UsageData {
            models,
            agents,
            daily,
            hourly,
            minutely,
            monthly,
            sessions,
            graph: Some(graph),
            total_tokens,
            total_cost,
            loading: false,
            error: None,
            current_streak,
            longest_streak,
        })
    }
}

fn parse_date(date_str: &str) -> Option<NaiveDate> {
    NaiveDate::parse_from_str(date_str, "%Y-%m-%d").ok()
}

/// Resolve a message to a Unix-ms timestamp, falling back to the `date`
/// string's midnight when `timestamp` is missing/zero. Local midnight is
/// preferred so the date renders on the same calendar day in the user's
/// timezone. DST edge cases are handled explicitly:
/// - Fall-back overlap (two valid midnights): take the earliest so the date
///   stays on the correct day.
/// - Spring-forward gap (midnight doesn't exist): fall back to UTC midnight
///   rather than silently returning 0 and losing the session boundary.
fn message_timestamp_ms(msg: &UnifiedMessage) -> i64 {
    message_timestamp_ms_with_timezone(msg, tokmesh_core::bucket_timezone())
}

fn message_timestamp_ms_with_timezone(
    msg: &UnifiedMessage,
    timezone: tokmesh_core::BucketTimezone,
) -> i64 {
    if msg.timestamp > 0 {
        return msg.timestamp;
    }
    parse_date(&msg.date)
        .map(|date| {
            timezone.start_of_day_ms(date).unwrap_or_else(|| {
                date.and_hms_opt(0, 0, 0)
                    .map(|wall| wall.and_utc().timestamp_millis())
                    .unwrap_or(0)
            })
        })
        .unwrap_or(0)
}

/// Convert Unix ms timestamp to a NaiveDateTime truncated to the pinned hour.
fn timestamp_to_hour(timestamp_ms: i64) -> Option<NaiveDateTime> {
    timestamp_to_hour_with_timezone(timestamp_ms, tokmesh_core::bucket_timezone())
}

fn timestamp_to_hour_with_timezone(
    timestamp_ms: i64,
    timezone: tokmesh_core::BucketTimezone,
) -> Option<NaiveDateTime> {
    if timestamp_ms <= 0 {
        return None;
    }
    let naive = timezone.naive_datetime_of_ms(timestamp_ms)?;
    Some(
        naive
            .date()
            .and_hms_opt(naive.hour(), 0, 0)
            .unwrap_or(naive),
    )
}

/// Derive an hour-truncated NaiveDateTime from `msg.timestamp` when present,
/// otherwise fall back to `msg.date`'s 00:00 bucket so messages with missing
/// timestamps are not silently dropped from hourly aggregation. Mirrors the
/// CLI hourly bucketing behavior in `tokmesh-core::lib::get_hourly_report`.
fn hour_bucket_with_fallback(timestamp_ms: i64, date_str: &str) -> Option<NaiveDateTime> {
    if let Some(dt) = timestamp_to_hour(timestamp_ms) {
        return Some(dt);
    }
    parse_date(date_str).and_then(|d| d.and_hms_opt(0, 0, 0))
}

/// Convert Unix ms timestamp to a NaiveDateTime truncated to the pinned minute.
fn timestamp_to_minute(timestamp_ms: i64) -> Option<NaiveDateTime> {
    timestamp_to_minute_with_timezone(timestamp_ms, tokmesh_core::bucket_timezone())
}

fn timestamp_to_minute_with_timezone(
    timestamp_ms: i64,
    timezone: tokmesh_core::BucketTimezone,
) -> Option<NaiveDateTime> {
    if timestamp_ms <= 0 {
        return None;
    }
    let naive = timezone.naive_datetime_of_ms(timestamp_ms)?;
    Some(
        naive
            .date()
            .and_hms_opt(naive.hour(), naive.minute(), 0)
            .unwrap_or(naive),
    )
}

/// Derive a minute-truncated NaiveDateTime from `msg.timestamp` when present,
/// otherwise fall back to `msg.date`'s 00:00 bucket so messages with missing
/// timestamps are not silently dropped from minutely aggregation.
fn minute_bucket_with_fallback(timestamp_ms: i64, date_str: &str) -> Option<NaiveDateTime> {
    if let Some(dt) = timestamp_to_minute(timestamp_ms) {
        return Some(dt);
    }
    parse_date(date_str).and_then(|d| d.and_hms_opt(0, 0, 0))
}

fn build_contribution_graph(daily: &[DailyUsage]) -> GraphData {
    build_contribution_graph_for_today(daily, tokmesh_core::bucket_timezone().today())
}

fn build_contribution_graph_for_today(daily: &[DailyUsage], today: NaiveDate) -> GraphData {
    if daily.is_empty() {
        return GraphData { weeks: vec![] };
    }

    let days_to_sunday = today.weekday().num_days_from_sunday();
    let end_date = today;
    let start_date = end_date - chrono::Duration::days(364 + days_to_sunday as i64);

    let daily_map: HashMap<NaiveDate, &DailyUsage> = daily.iter().map(|d| (d.date, d)).collect();

    let max_cost = daily.iter().map(|d| d.cost).fold(0.0_f64, |a, b| a.max(b));

    let mut weeks: Vec<Vec<Option<ContributionDay>>> = Vec::new();
    let mut current_week: Vec<Option<ContributionDay>> = Vec::new();

    let mut current_date = start_date;
    while current_date <= end_date {
        let day = if let Some(usage) = daily_map.get(&current_date) {
            let raw_intensity = if max_cost > 0.0 {
                usage.cost / max_cost
            } else {
                0.0
            };
            let intensity = if raw_intensity.is_finite() {
                raw_intensity.clamp(0.0, 1.0)
            } else {
                0.0
            };
            Some(ContributionDay {
                date: current_date,
                tokens: usage.tokens.total(),
                cost: usage.cost,
                intensity,
            })
        } else {
            Some(ContributionDay {
                date: current_date,
                tokens: 0,
                cost: 0.0,
                intensity: 0.0,
            })
        };

        current_week.push(day);

        if current_date.weekday() == chrono::Weekday::Sat || current_date == end_date {
            weeks.push(current_week);
            current_week = Vec::new();
        }

        current_date += chrono::Duration::days(1);
    }

    GraphData { weeks }
}

fn calculate_streaks(daily: &[DailyUsage]) -> (u32, u32) {
    calculate_streaks_for_today(daily, tokmesh_core::bucket_timezone().today())
}

pub fn aggregate_monthly_from_daily(daily: &[DailyUsage]) -> Vec<MonthlyUsage> {
    let mut monthly_map: HashMap<String, MonthlyUsage> = HashMap::new();

    for day in daily {
        let month = day.date.format("%Y-%m").to_string();
        let entry = monthly_map
            .entry(month.clone())
            .or_insert_with(|| MonthlyUsage {
                month,
                tokens: TokenBreakdown::default(),
                cost: 0.0,
                message_count: 0,
                turn_count: 0,
            });

        entry.tokens.input = entry.tokens.input.saturating_add(day.tokens.input);
        entry.tokens.output = entry.tokens.output.saturating_add(day.tokens.output);
        entry.tokens.cache_read = entry
            .tokens
            .cache_read
            .saturating_add(day.tokens.cache_read);
        entry.tokens.cache_write = entry
            .tokens
            .cache_write
            .saturating_add(day.tokens.cache_write);
        entry.tokens.reasoning = entry.tokens.reasoning.saturating_add(day.tokens.reasoning);
        entry.cost += day.cost;
        entry.message_count = entry.message_count.saturating_add(day.message_count);
        entry.turn_count = entry.turn_count.saturating_add(day.turn_count);
    }

    let mut monthly: Vec<MonthlyUsage> = monthly_map.into_values().collect();
    monthly.sort_by(|a, b| b.month.cmp(&a.month));
    monthly
}

fn calculate_streaks_for_today(daily: &[DailyUsage], today: NaiveDate) -> (u32, u32) {
    if daily.is_empty() {
        return (0, 0);
    }

    let dates: HashSet<NaiveDate> = daily.iter().map(|d| d.date).collect();

    let mut current_streak = 0u32;
    let mut check_date = today;

    while dates.contains(&check_date) {
        current_streak += 1;
        check_date -= chrono::Duration::days(1);
    }

    if current_streak == 0 {
        let yesterday = today - chrono::Duration::days(1);
        check_date = yesterday;
        while dates.contains(&check_date) {
            current_streak += 1;
            check_date -= chrono::Duration::days(1);
        }
    }

    let mut longest_streak = 0u32;
    let mut sorted_dates: Vec<NaiveDate> = dates.into_iter().collect();
    sorted_dates.sort();

    let mut streak = 0u32;
    let mut prev_date: Option<NaiveDate> = None;

    for date in sorted_dates {
        if let Some(prev) = prev_date {
            if date == prev + chrono::Duration::days(1) {
                streak += 1;
            } else {
                longest_streak = longest_streak.max(streak);
                streak = 1;
            }
        } else {
            streak = 1;
        }
        prev_date = Some(date);
    }
    longest_streak = longest_streak.max(streak);

    (current_streak, longest_streak)
}

/// Time-of-day period bucket for profile view
#[derive(Debug, Clone)]
pub struct PeriodBucket {
    pub label: &'static str,
    pub hour_range: &'static str,
    pub total_tokens: u64,
}

/// Weekday bucket for profile view
#[derive(Debug, Clone)]
pub struct WeekdayBucket {
    pub day: &'static str,
    pub total_tokens: u64,
}

/// Aggregate hourly data into time-of-day periods
pub fn aggregate_by_period(hourly: &[HourlyUsage]) -> Vec<PeriodBucket> {
    let periods: [(&str, &str, Vec<usize>); 4] = [
        ("Morning", "05:00-11:59", (5..=11).collect()),
        ("Daytime", "12:00-16:59", (12..=16).collect()),
        ("Evening", "17:00-21:59", (17..=21).collect()),
        ("Night", "22:00-04:59", vec![22, 23, 0, 1, 2, 3, 4]),
    ];

    periods
        .iter()
        .map(|(label, hour_range, hours)| {
            let mut total_tokens = 0u64;

            for entry in hourly {
                let hour = entry.datetime.hour() as usize;
                if hours.contains(&hour) {
                    total_tokens = total_tokens.saturating_add(entry.tokens.total());
                }
            }

            PeriodBucket {
                label,
                hour_range,
                total_tokens,
            }
        })
        .collect()
}

/// Aggregate hourly data by weekday
pub fn aggregate_by_weekday(hourly: &[HourlyUsage]) -> Vec<WeekdayBucket> {
    use chrono::Datelike;

    let weekdays = [
        "Monday",
        "Tuesday",
        "Wednesday",
        "Thursday",
        "Friday",
        "Saturday",
        "Sunday",
    ];
    let mut buckets: Vec<u64> = vec![0; 7];

    for entry in hourly {
        let weekday = entry.datetime.weekday().num_days_from_monday() as usize;
        buckets[weekday] = buckets[weekday].saturating_add(entry.tokens.total());
    }

    weekdays
        .iter()
        .enumerate()
        .map(|(i, day)| WeekdayBucket {
            day,
            total_tokens: buckets[i],
        })
        .collect()
}

/// Find peak hour across all hourly data
pub fn find_peak_hour(hourly: &[HourlyUsage]) -> Option<(u32, u64, f64)> {
    use std::collections::HashMap;

    let mut hour_totals: HashMap<u32, (u64, f64)> = HashMap::new();

    for entry in hourly {
        let hour = entry.datetime.hour();
        let entry_totals = hour_totals.entry(hour).or_insert((0, 0.0));
        entry_totals.0 = entry_totals.0.saturating_add(entry.tokens.total());
        entry_totals.1 += entry.cost;
    }

    hour_totals
        .into_iter()
        .max_by_key(|(_, (tokens, _))| *tokens)
        .map(|(hour, (tokens, cost))| (hour, tokens, cost))
}

#[cfg(test)]
mod tests {
    use super::*;
    use serial_test::serial;
    use std::collections::HashMap;
    use std::env;
    use std::fs;
    use tempfile::TempDir;
    use tokio::runtime::{Handle, Runtime};
    use tokmesh_core::parse_local_unified_messages_with_pricing;
    use tokmesh_core::pricing::{ModelPricing, PricingService};
    use tokmesh_core::TokenBreakdown as CoreTokenBreakdown;

    #[test]
    fn positive_unified_token_total_saturates_instead_of_overflowing() {
        // tokmesh-core clamps corrupt per-field token buckets to i64::MAX; a
        // plain `+` fold over two clamped buckets would panic in debug builds.
        let tokens = CoreTokenBreakdown {
            input: i64::MAX,
            output: 0,
            cache_read: i64::MAX,
            cache_write: -5,
            reasoning: 0,
        };
        assert_eq!(positive_unified_token_total(&tokens), i64::MAX);
    }

    #[test]
    fn hourly_and_minutely_buckets_use_pinned_timezone() {
        use chrono::{TimeZone, Utc};

        let timezone = tokmesh_core::parse_bucket_timezone("Pacific/Kiritimati").unwrap();
        let timestamp = Utc
            .with_ymd_and_hms(2026, 7, 31, 12, 34, 56)
            .single()
            .unwrap()
            .timestamp_millis();

        assert_eq!(
            timestamp_to_hour_with_timezone(timestamp, timezone)
                .unwrap()
                .format("%Y-%m-%d %H:%M")
                .to_string(),
            "2026-08-01 02:00"
        );
        assert_eq!(
            timestamp_to_minute_with_timezone(timestamp, timezone)
                .unwrap()
                .format("%Y-%m-%d %H:%M")
                .to_string(),
            "2026-08-01 02:34"
        );
    }

    #[test]
    fn missing_timestamp_uses_pinned_date_midnight() {
        use chrono::{TimeZone, Utc};

        let timezone = tokmesh_core::parse_bucket_timezone("Pacific/Kiritimati").unwrap();
        let mut message = UnifiedMessage::new(
            "opencode",
            "gpt-4o",
            "openai",
            "session",
            0,
            CoreTokenBreakdown::default(),
            0.0,
        );
        message.date = "2026-08-01".to_string();

        assert_eq!(
            message_timestamp_ms_with_timezone(&message, timezone),
            Utc.with_ymd_and_hms(2026, 7, 31, 10, 0, 0)
                .single()
                .unwrap()
                .timestamp_millis()
        );
    }

    fn test_pricing_service() -> PricingService {
        let mut litellm = HashMap::new();
        litellm.insert(
            "claude-sonnet-4".into(),
            ModelPricing {
                input_cost_per_token: Some(0.00001),
                output_cost_per_token: Some(0.00002),
                cache_read_input_token_cost: Some(0.000003),
                ..Default::default()
            },
        );
        litellm.insert(
            "claude-haiku-4".into(),
            ModelPricing {
                input_cost_per_token: Some(0.000004),
                output_cost_per_token: Some(0.000006),
                cache_read_input_token_cost: Some(0.000001),
                ..Default::default()
            },
        );
        litellm.insert(
            "accounts/fireworks/models/deepseek-v3-0324".into(),
            ModelPricing {
                input_cost_per_token: Some(0.01),
                output_cost_per_token: Some(0.03),
                ..Default::default()
            },
        );

        PricingService::new(litellm, HashMap::new())
    }

    fn load_with_pricing(
        loader: &DataLoader,
        enabled_clients: &[ClientId],
        group_by: &GroupBy,
        include_synthetic: bool,
        pricing: Option<&PricingService>,
    ) -> Result<UsageData> {
        let home = dirs::home_dir()
            .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))?
            .to_string_lossy()
            .to_string();

        let mut sources: Vec<String> = enabled_clients
            .iter()
            .map(|client| client.as_str().to_string())
            .collect();
        if include_synthetic {
            sources.push("synthetic".to_string());
        }

        let opts = LocalParseOptions {
            home_dir: Some(home),
            use_env_roots: true,
            clients: Some(sources),
            since: loader.since.clone(),
            until: loader.until.clone(),
            year: loader.year.clone(),
            scanner_settings: data_loader_scanner_settings(),
        };

        let messages = if Handle::try_current().is_ok() {
            std::thread::scope(|s| {
                s.spawn(|| {
                    let rt = Runtime::new().map_err(|e| e.to_string())?;
                    rt.block_on(parse_local_unified_messages_with_pricing(opts, pricing))
                })
                .join()
                .unwrap_or_else(|_| Err("data loader thread panicked".to_string()))
            })
        } else {
            Runtime::new()?.block_on(parse_local_unified_messages_with_pricing(opts, pricing))
        }
        .map_err(anyhow::Error::msg)?;

        loader.aggregate_messages(messages, group_by)
    }

    fn expected_message_cost(
        pricing: &PricingService,
        model_id: &str,
        provider_id: &str,
        tokens: CoreTokenBreakdown,
    ) -> f64 {
        pricing.calculate_cost_with_provider(model_id, Some(provider_id), &tokens)
    }

    fn assert_cost_matches(actual: f64, expected: f64) {
        assert!(
            (actual - expected).abs() < 1e-9,
            "expected cost {expected}, got {actual}"
        );
    }

    fn make_workspace_message(
        client: &str,
        model_id: &str,
        provider_id: &str,
        session_id: &str,
        cost: f64,
        workspace_key: Option<&str>,
        workspace_label: Option<&str>,
    ) -> UnifiedMessage {
        let mut msg = UnifiedMessage::new(
            client,
            model_id,
            provider_id,
            session_id,
            1_735_689_600_000,
            tokmesh_core::TokenBreakdown {
                input: 10,
                output: 5,
                cache_read: 0,
                cache_write: 0,
                reasoning: 0,
            },
            cost,
        );
        msg.set_workspace(
            workspace_key.map(str::to_string),
            workspace_label.map(str::to_string),
        );
        msg
    }

    #[test]
    fn test_client_all() {
        let clients = ClientId::ALL;
        assert_eq!(clients.len(), ClientId::COUNT);
        assert_eq!(clients[0], ClientId::OpenCode);
        assert_eq!(clients[1], ClientId::Claude);
        assert_eq!(clients[2], ClientId::Codex);
        assert_eq!(clients[3], ClientId::Cursor);
        assert_eq!(clients[4], ClientId::Gemini);
        assert_eq!(clients[5], ClientId::Amp);
        assert_eq!(clients[6], ClientId::Droid);
        assert_eq!(clients[7], ClientId::OpenClaw);
        assert_eq!(clients[8], ClientId::Pi);
        assert_eq!(clients[9], ClientId::Kimi);
        assert_eq!(clients[10], ClientId::Qwen);
        assert_eq!(clients[11], ClientId::RooCode);
        assert_eq!(clients[12], ClientId::KiloCode);
        assert_eq!(clients[13], ClientId::Mux);
        assert_eq!(clients[14], ClientId::Kilo);
        assert_eq!(clients[15], ClientId::Crush);
        assert_eq!(clients[16], ClientId::Hermes);
        assert_eq!(clients[17], ClientId::Copilot);
        assert_eq!(clients[18], ClientId::Goose);
        assert_eq!(clients[19], ClientId::Codebuff);
        assert_eq!(clients[20], ClientId::Antigravity);
        assert_eq!(clients[21], ClientId::Zed);
        assert_eq!(clients[22], ClientId::Kiro);
        assert_eq!(clients[23], ClientId::Trae);
        assert_eq!(clients[24], ClientId::Warp);
        assert_eq!(clients[25], ClientId::Cline);
        assert_eq!(clients[26], ClientId::Gjc);
        assert_eq!(clients[27], ClientId::Grok);
        assert_eq!(clients[28], ClientId::Jcode);
        assert_eq!(clients[29], ClientId::CommandCode);
        assert_eq!(clients[30], ClientId::MiMoCode);
        assert_eq!(clients[31], ClientId::AntigravityCli);
        assert_eq!(clients[32], ClientId::Junie);
        assert_eq!(clients[33], ClientId::Zcode);
        assert_eq!(clients[34], ClientId::OpenCodeReview);
        assert_eq!(clients[35], ClientId::CodeBuddy);
        assert_eq!(clients[36], ClientId::WorkBuddy);
    }

    #[test]
    fn test_client_as_str() {
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::OpenCode),
            "OpenCode"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Claude),
            "Claude"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Codex),
            "Codex"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Copilot),
            "Copilot"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Cursor),
            "Cursor"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Gemini),
            "Gemini"
        );
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Amp), "Amp");
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Droid),
            "Droid"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::OpenClaw),
            "OpenClaw"
        );
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Pi), "Pi");
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Kimi), "Kimi");
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Qwen), "Qwen");
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::RooCode),
            "Roo Code"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::KiloCode),
            "KiloCode"
        );
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Mux), "Mux");
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Kilo),
            "Kilo CLI"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Crush),
            "Crush"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Hermes),
            "Hermes Agent"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Codebuff),
            "Codebuff"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Antigravity),
            "Antigravity"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Zed),
            "Zed Agent"
        );
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Kiro), "Kiro");
        assert_eq!(crate::tui::client_ui::display_name(ClientId::Trae), "Trae");
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Cline),
            "Cline"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Grok),
            "Grok Build"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Jcode),
            "Jcode"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::AntigravityCli),
            "Antigravity CLI"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::Junie),
            "Junie"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::CodeBuddy),
            "CodeBuddy"
        );
        assert_eq!(
            crate::tui::client_ui::display_name(ClientId::WorkBuddy),
            "WorkBuddy"
        );
    }

    #[test]
    fn test_client_key() {
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::OpenCode), '1');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Claude), '2');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Codex), '3');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Copilot), 'c');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Cursor), '4');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Gemini), '5');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Amp), '6');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Droid), '7');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::OpenClaw), '8');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Pi), '9');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Kimi), '0');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Qwen), 'w');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::RooCode), 'r');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::KiloCode), 'k');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Mux), 'x');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Kilo), 'l');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Crush), 'h');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Hermes), 'e');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Codebuff), 'b');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Antigravity), 'a');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Zed), 'z');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Kiro), 'i');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Trae), 'y');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Cline), 'n');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Gjc), 'g');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Grok), 'u');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Jcode), 'j');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::AntigravityCli), 'f');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::Junie), 'p');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::CodeBuddy), 'C');
        assert_eq!(crate::tui::client_ui::hotkey(ClientId::WorkBuddy), 'B');
    }

    #[test]
    fn test_client_from_key() {
        assert_eq!(
            crate::tui::client_ui::from_hotkey('1'),
            Some(ClientId::OpenCode)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('2'),
            Some(ClientId::Claude)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('3'),
            Some(ClientId::Codex)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('c'),
            Some(ClientId::Copilot)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('4'),
            Some(ClientId::Cursor)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('5'),
            Some(ClientId::Gemini)
        );
        assert_eq!(crate::tui::client_ui::from_hotkey('6'), Some(ClientId::Amp));
        assert_eq!(
            crate::tui::client_ui::from_hotkey('7'),
            Some(ClientId::Droid)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('8'),
            Some(ClientId::OpenClaw)
        );
        assert_eq!(crate::tui::client_ui::from_hotkey('9'), Some(ClientId::Pi));
        assert_eq!(
            crate::tui::client_ui::from_hotkey('0'),
            Some(ClientId::Kimi)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('w'),
            Some(ClientId::Qwen)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('r'),
            Some(ClientId::RooCode)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('k'),
            Some(ClientId::KiloCode)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('l'),
            Some(ClientId::Kilo)
        );
        assert_eq!(crate::tui::client_ui::from_hotkey('x'), Some(ClientId::Mux));
        assert_eq!(
            crate::tui::client_ui::from_hotkey('h'),
            Some(ClientId::Crush)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('e'),
            Some(ClientId::Hermes)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('b'),
            Some(ClientId::Codebuff)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('a'),
            Some(ClientId::Antigravity)
        );
        assert_eq!(crate::tui::client_ui::from_hotkey('z'), Some(ClientId::Zed));
        assert_eq!(
            crate::tui::client_ui::from_hotkey('i'),
            Some(ClientId::Kiro)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('y'),
            Some(ClientId::Trae)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('u'),
            Some(ClientId::Grok)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('j'),
            Some(ClientId::Jcode)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('f'),
            Some(ClientId::AntigravityCli)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('p'),
            Some(ClientId::Junie)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('C'),
            Some(ClientId::CodeBuddy)
        );
        assert_eq!(
            crate::tui::client_ui::from_hotkey('B'),
            Some(ClientId::WorkBuddy)
        );
    }

    #[test]
    fn test_token_breakdown_total() {
        let breakdown = TokenBreakdown {
            input: 100,
            output: 200,
            cache_read: 50,
            cache_write: 25,
            reasoning: 10,
        };
        assert_eq!(breakdown.total(), 385);
    }

    #[test]
    fn test_token_breakdown_total_with_overflow() {
        let breakdown = TokenBreakdown {
            input: u64::MAX,
            output: 1,
            cache_read: 0,
            cache_write: 0,
            reasoning: 0,
        };
        // saturating_add should prevent overflow
        assert_eq!(breakdown.total(), u64::MAX);
    }

    #[test]
    fn test_token_breakdown_default() {
        let breakdown = TokenBreakdown::default();
        assert_eq!(breakdown.input, 0);
        assert_eq!(breakdown.output, 0);
        assert_eq!(breakdown.cache_read, 0);
        assert_eq!(breakdown.cache_write, 0);
        assert_eq!(breakdown.reasoning, 0);
        assert_eq!(breakdown.total(), 0);
    }

    #[test]
    fn test_data_loader_new() {
        let loader = DataLoader::new(None);
        assert!(loader._sessions_path.is_none());
        assert!(loader.since.is_none());
        assert!(loader.until.is_none());
        assert!(loader.year.is_none());
    }

    #[test]
    fn test_data_loader_scanner_settings_is_hermetic_under_cfg_test() {
        // Regression guard: the `#[cfg(test)]` branch of
        // `data_loader_scanner_settings` must not read
        // `~/.config/tokmesh/settings.json`. Otherwise every DataLoader
        // unit test becomes machine-dependent as soon as a developer
        // pins extra OpenCode dbs in their real settings.json.
        //
        // This test cannot sandbox HOME (many of the sibling tests in
        // this module would race against each other if it did), so
        // instead it asserts the cfg(test) helper returns a default
        // ScannerSettings regardless of what the real settings file
        // contains on the developer's machine.
        let settings = super::data_loader_scanner_settings();
        assert!(
            settings.opencode_db_paths.is_empty(),
            "under #[cfg(test)] data_loader_scanner_settings must return \
             ScannerSettings::default() so unit tests stay hermetic, but \
             got {:?}",
            settings.opencode_db_paths
        );
    }

    #[test]
    fn test_data_loader_with_filters() {
        let loader = DataLoader::with_filters(
            Some(PathBuf::from("/tmp/sessions")),
            Some("2024-01-01".to_string()),
            Some("2024-12-31".to_string()),
            Some("2024".to_string()),
        );

        assert_eq!(loader._sessions_path, Some(PathBuf::from("/tmp/sessions")));
        assert_eq!(loader.since, Some("2024-01-01".to_string()));
        assert_eq!(loader.until, Some("2024-12-31".to_string()));
        assert_eq!(loader.year, Some("2024".to_string()));
    }

    #[test]
    fn test_parse_date() {
        assert_eq!(
            parse_date("2024-01-15"),
            Some(NaiveDate::from_ymd_opt(2024, 1, 15).unwrap())
        );
        assert_eq!(
            parse_date("2024-12-31"),
            Some(NaiveDate::from_ymd_opt(2024, 12, 31).unwrap())
        );
        assert_eq!(parse_date("invalid"), None);
        assert_eq!(parse_date("2024-13-01"), None);
        assert_eq!(parse_date(""), None);
    }

    #[test]
    fn test_build_contribution_graph_uses_provided_today() {
        let today = NaiveDate::from_ymd_opt(2026, 3, 8).unwrap();
        let graph = build_contribution_graph_for_today(&[], today);
        assert!(graph.weeks.is_empty());

        let daily = vec![DailyUsage {
            date: NaiveDate::from_ymd_opt(2026, 3, 2).unwrap(),
            tokens: TokenBreakdown::default(),
            cost: 0.0,
            source_breakdown: BTreeMap::new(),
            message_count: 0,
            turn_count: 0,
        }];
        let graph = build_contribution_graph_for_today(&daily, today);
        let last_day = graph
            .weeks
            .last()
            .and_then(|week| week.last())
            .and_then(|day| day.as_ref())
            .map(|day| day.date);
        assert_eq!(last_day, Some(today));
    }

    #[test]
    fn test_aggregate_messages_builds_agent_usage() {
        let loader = DataLoader::new(None);
        let messages = vec![
            UnifiedMessage::new_with_agent(
                "opencode",
                "claude-sonnet-4",
                "anthropic",
                "session-1",
                1_735_689_600_000,
                tokmesh_core::TokenBreakdown {
                    input: 10,
                    output: 5,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                1.25,
                Some("builder".to_string()),
            ),
            UnifiedMessage::new_with_agent(
                "roocode",
                "claude-sonnet-4",
                "anthropic",
                "session-2",
                1_735_689_700_000,
                tokmesh_core::TokenBreakdown {
                    input: 20,
                    output: 10,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                2.75,
                Some("builder".to_string()),
            ),
        ];

        let usage = loader
            .aggregate_messages(messages, &GroupBy::Model)
            .unwrap();

        assert_eq!(usage.agents.len(), 1);
        assert_eq!(usage.agents[0].agent, "Builder");
        assert_eq!(usage.agents[0].clients, "opencode, roocode");
        assert_eq!(usage.agents[0].message_count, 2);
        assert!((usage.agents[0].cost - 4.0).abs() < f64::EPSILON);
        assert_eq!(usage.agents[0].tokens.total(), 45);
    }

    #[test]
    fn total_tokens_saturates_across_corrupt_model_buckets() {
        let loader = DataLoader::new(None);
        // Three distinct models each carrying i64::MAX input tokens: no
        // single model bucket overflows (TokenBreakdown::total saturates
        // internally), but summing three of them plainly overflows u64.
        let messages: Vec<UnifiedMessage> = (0..3)
            .map(|i| {
                UnifiedMessage::new(
                    "claude",
                    format!("model-{i}"),
                    "anthropic",
                    format!("session-{i}"),
                    1_735_689_600_000,
                    tokmesh_core::TokenBreakdown {
                        input: i64::MAX,
                        output: 0,
                        cache_read: 0,
                        cache_write: 0,
                        reasoning: 0,
                    },
                    0.0,
                )
            })
            .collect();

        let usage = loader
            .aggregate_messages(messages, &GroupBy::Model)
            .unwrap();

        assert_eq!(usage.models.len(), 3);
        assert_eq!(usage.total_tokens, u64::MAX);
    }

    #[test]
    fn test_aggregate_messages_groups_by_workspace_and_model() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1.25,
                        Some("/repo-a"),
                        Some("repo-a"),
                    ),
                    make_workspace_message(
                        "qwen",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        2.75,
                        Some("/repo-a"),
                        Some("repo-a"),
                    ),
                ],
                &GroupBy::WorkspaceModel,
            )
            .unwrap();

        assert_eq!(usage.models.len(), 1);
        assert_eq!(usage.models[0].workspace_key.as_deref(), Some("/repo-a"));
        assert_eq!(usage.models[0].workspace_label.as_deref(), Some("repo-a"));
        assert_eq!(usage.models[0].model, "claude-sonnet-4-5");
        assert_eq!(usage.models[0].client, "claude, qwen");
        assert_eq!(usage.models[0].session_count, 2);
        assert_eq!(usage.models[0].cost, 4.0);
    }

    #[test]
    fn test_aggregate_messages_workspace_grouping_keeps_unknown_bucket_visible() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1.0,
                        None,
                        None,
                    ),
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        2.0,
                        None,
                        None,
                    ),
                ],
                &GroupBy::WorkspaceModel,
            )
            .unwrap();

        assert_eq!(usage.models.len(), 1);
        assert_eq!(usage.models[0].workspace_key, None);
        assert_eq!(
            usage.models[0].workspace_label.as_deref(),
            Some(UNKNOWN_WORKSPACE_LABEL)
        );
        assert_eq!(usage.models[0].session_count, 2);
        assert_eq!(usage.models[0].cost, 3.0);
    }

    #[test]
    fn test_aggregate_messages_workspace_grouping_keeps_real_unknown_workspace_separate() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1.0,
                        Some("unknown-workspace"),
                        Some("unknown-workspace"),
                    ),
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        2.0,
                        None,
                        None,
                    ),
                ],
                &GroupBy::WorkspaceModel,
            )
            .unwrap();

        assert_eq!(usage.models.len(), 2);
        assert!(usage.models.iter().any(|model| {
            model.workspace_key.as_deref() == Some("unknown-workspace")
                && model.workspace_label.as_deref() == Some("unknown-workspace")
                && (model.cost - 1.0).abs() < f64::EPSILON
        }));
        assert!(usage.models.iter().any(|model| {
            model.workspace_key.is_none()
                && model.workspace_label.as_deref() == Some(UNKNOWN_WORKSPACE_LABEL)
                && (model.cost - 2.0).abs() < f64::EPSILON
        }));
    }

    #[test]
    fn test_aggregate_messages_workspace_grouping_splits_daily_models_by_workspace() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1.0,
                        Some("/repo-a"),
                        Some("repo-a"),
                    ),
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        2.0,
                        Some("/repo-b"),
                        Some("repo-b"),
                    ),
                ],
                &GroupBy::WorkspaceModel,
            )
            .unwrap();

        assert_eq!(usage.daily.len(), 1);
        let claude = usage.daily[0].source_breakdown.get("claude").unwrap();
        let daily_keys: Vec<_> = claude.models.keys().cloned().collect();
        assert_eq!(daily_keys.len(), 2);
        assert_ne!(daily_keys[0], daily_keys[1]);
        let daily_display_names: Vec<_> = claude
            .models
            .values()
            .map(|info| info.display_name.clone())
            .collect();
        assert_eq!(
            daily_display_names,
            vec![
                "repo-a / claude-sonnet-4-5".to_string(),
                "repo-b / claude-sonnet-4-5".to_string()
            ]
        );
    }

    #[test]
    fn test_aggregate_messages_workspace_grouping_disambiguates_identical_labels() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1.0,
                        Some("/srv/team-a/demo"),
                        Some("demo"),
                    ),
                    make_workspace_message(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        2.0,
                        Some("/srv/team-b/demo"),
                        Some("demo"),
                    ),
                ],
                &GroupBy::WorkspaceModel,
            )
            .unwrap();

        assert_eq!(usage.daily.len(), 1);
        let claude = usage.daily[0].source_breakdown.get("claude").unwrap();
        assert_eq!(claude.models.len(), 2);

        // Keys must differ even though display names are identical
        let daily_keys: Vec<_> = claude.models.keys().cloned().collect();
        assert_eq!(daily_keys.len(), 2);
        assert_ne!(daily_keys[0], daily_keys[1]);

        let display_names: Vec<_> = claude
            .models
            .values()
            .map(|info| info.display_name.clone())
            .collect();
        assert_eq!(
            display_names,
            vec![
                "demo / claude-sonnet-4-5".to_string(),
                "demo / claude-sonnet-4-5".to_string()
            ]
        );
    }

    #[test]
    fn test_aggregate_messages_client_provider_model_splits_providers_in_daily_breakdown() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    UnifiedMessage::new(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1_735_689_600_000,
                        tokmesh_core::TokenBreakdown {
                            input: 10,
                            output: 5,
                            cache_read: 0,
                            cache_write: 0,
                            reasoning: 0,
                        },
                        1.0,
                    ),
                    UnifiedMessage::new(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "github-copilot",
                        "session-2",
                        1_735_689_600_000,
                        tokmesh_core::TokenBreakdown {
                            input: 20,
                            output: 10,
                            cache_read: 0,
                            cache_write: 0,
                            reasoning: 0,
                        },
                        2.0,
                    ),
                ],
                &GroupBy::ClientProviderModel,
            )
            .unwrap();

        assert_eq!(usage.daily.len(), 1);
        let claude = usage.daily[0].source_breakdown.get("claude").unwrap();
        assert_eq!(claude.models.len(), 2);

        let anthropic_key = "anthropic:claude-sonnet-4-5";
        let copilot_key = "github-copilot:claude-sonnet-4-5";
        let anthropic_model = claude.models.get(anthropic_key).unwrap();
        assert_eq!(
            anthropic_model.display_name,
            "anthropic / claude-sonnet-4-5"
        );
        assert_eq!(anthropic_model.provider, "anthropic");
        assert_eq!(anthropic_model.tokens.total(), 15);
        assert_eq!(anthropic_model.messages, 1);

        let copilot_model = claude.models.get(copilot_key).unwrap();
        assert_eq!(
            copilot_model.display_name,
            "github-copilot / claude-sonnet-4-5"
        );
        assert_eq!(copilot_model.provider, "github-copilot");
        assert_eq!(copilot_model.tokens.total(), 30);
        assert_eq!(copilot_model.messages, 1);
    }

    #[test]
    fn test_aggregate_messages_keeps_same_model_split_across_sources_in_daily_breakdown() {
        let loader = DataLoader::new(None);
        let usage = loader
            .aggregate_messages(
                vec![
                    UnifiedMessage::new(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        1_735_689_600_000,
                        tokmesh_core::TokenBreakdown {
                            input: 10,
                            output: 5,
                            cache_read: 0,
                            cache_write: 0,
                            reasoning: 0,
                        },
                        1.0,
                    ),
                    UnifiedMessage::new(
                        "cursor",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-2",
                        1_735_689_600_000,
                        tokmesh_core::TokenBreakdown {
                            input: 20,
                            output: 10,
                            cache_read: 0,
                            cache_write: 0,
                            reasoning: 0,
                        },
                        2.0,
                    ),
                ],
                &GroupBy::Model,
            )
            .unwrap();

        assert_eq!(usage.daily.len(), 1);
        assert_eq!(usage.daily[0].source_breakdown.len(), 2);

        let claude = usage.daily[0].source_breakdown.get("claude").unwrap();
        assert_eq!(claude.cost, 1.0);
        assert_eq!(claude.models.len(), 1);
        let claude_model = claude.models.get("claude-sonnet-4-5").unwrap();
        assert_eq!(claude_model.display_name, "claude-sonnet-4-5");
        assert_eq!(claude_model.tokens.total(), 15);

        let cursor = usage.daily[0].source_breakdown.get("cursor").unwrap();
        assert_eq!(cursor.cost, 2.0);
        assert_eq!(cursor.models.len(), 1);
        let cursor_model = cursor.models.get("claude-sonnet-4-5").unwrap();
        assert_eq!(cursor_model.display_name, "claude-sonnet-4-5");
        assert_eq!(cursor_model.tokens.total(), 30);
    }

    #[test]
    fn test_aggregate_messages_merges_oh_my_opencode_agent_variants() {
        let loader = DataLoader::new(None);
        let messages = vec![
            UnifiedMessage::new_with_agent(
                "opencode",
                "claude-opus-4-6",
                "anthropic",
                "session-1",
                1_735_689_600_000,
                tokmesh_core::TokenBreakdown {
                    input: 10,
                    output: 5,
                    cache_read: 100,
                    cache_write: 20,
                    reasoning: 0,
                },
                1.5,
                Some("Sisyphus".to_string()),
            ),
            UnifiedMessage::new_with_agent(
                "opencode",
                "claude-opus-4-6",
                "anthropic",
                "session-2",
                1_735_689_700_000,
                tokmesh_core::TokenBreakdown {
                    input: 20,
                    output: 10,
                    cache_read: 200,
                    cache_write: 40,
                    reasoning: 0,
                },
                2.5,
                Some("Sisyphus (Ultraworker)".to_string()),
            ),
        ];

        let usage = loader
            .aggregate_messages(messages, &GroupBy::Model)
            .unwrap();

        assert_eq!(usage.agents.len(), 1);
        assert_eq!(usage.agents[0].agent, "Sisyphus");
        assert_eq!(usage.agents[0].clients, "opencode");
        assert_eq!(usage.agents[0].message_count, 2);
        assert!((usage.agents[0].cost - 4.0).abs() < f64::EPSILON);
        assert_eq!(usage.agents[0].tokens.total(), 405);
    }

    #[test]
    fn test_aggregate_messages_merges_opencode_agent_case_variants() {
        let loader = DataLoader::new(None);
        let messages = vec![
            UnifiedMessage::new_with_agent(
                "opencode",
                "claude-opus-4-6",
                "anthropic",
                "session-1",
                1_735_689_600_000,
                tokmesh_core::TokenBreakdown {
                    input: 10,
                    output: 5,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                1.5,
                Some("Hephaestus".to_string()),
            ),
            UnifiedMessage::new_with_agent(
                "opencode",
                "claude-opus-4-6",
                "anthropic",
                "session-2",
                1_735_689_700_000,
                tokmesh_core::TokenBreakdown {
                    input: 20,
                    output: 10,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                2.5,
                Some("hephaestus".to_string()),
            ),
        ];

        let usage = loader
            .aggregate_messages(messages, &GroupBy::Model)
            .unwrap();

        assert_eq!(usage.agents.len(), 1);
        assert_eq!(usage.agents[0].agent, "Hephaestus");
        assert_eq!(usage.agents[0].clients, "opencode");
        assert_eq!(usage.agents[0].message_count, 2);
        assert!((usage.agents[0].cost - 4.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_aggregate_messages_does_not_merge_omo_variants_for_non_opencode_clients() {
        let loader = DataLoader::new(None);
        let messages = vec![
            UnifiedMessage::new_with_agent(
                "claude",
                "claude-opus-4-6",
                "anthropic",
                "session-1",
                1_735_689_600_000,
                tokmesh_core::TokenBreakdown {
                    input: 10,
                    output: 5,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                1.5,
                Some("Sisyphus".to_string()),
            ),
            UnifiedMessage::new_with_agent(
                "claude",
                "claude-opus-4-6",
                "anthropic",
                "session-2",
                1_735_689_700_000,
                tokmesh_core::TokenBreakdown {
                    input: 20,
                    output: 10,
                    cache_read: 0,
                    cache_write: 0,
                    reasoning: 0,
                },
                2.5,
                Some("Sisyphus (Ultraworker)".to_string()),
            ),
        ];

        let usage = loader
            .aggregate_messages(messages, &GroupBy::Model)
            .unwrap();

        assert_eq!(usage.agents.len(), 2);
        assert!(usage.agents.iter().any(|agent| agent.agent == "Sisyphus"));
        assert!(usage
            .agents
            .iter()
            .any(|agent| agent.agent == "Sisyphus (Ultraworker)"));
    }

    #[test]
    #[serial]
    fn test_data_loader_loads_agent_usage_from_roocode_files() {
        let temp_dir = TempDir::new().unwrap();
        let previous_home = env::var_os("HOME");
        let task_root = temp_dir
            .path()
            .join(".config/Code/User/globalStorage/rooveterinaryinc.roo-cline/tasks");

        let architect_dir = task_root.join("task-architect");
        fs::create_dir_all(&architect_dir).unwrap();
        fs::write(
            architect_dir.join("ui_messages.json"),
            r#"[
  {
    "type": "say",
    "say": "api_req_started",
    "ts": "2026-03-07T16:00:00Z",
    "text": "{\"cost\":8.4,\"tokensIn\":420000,\"tokensOut\":120000,\"cacheReads\":32000,\"cacheWrites\":0,\"apiProtocol\":\"anthropic\"}"
  },
  {
    "type": "say",
    "say": "api_req_started",
    "ts": "2026-03-07T16:05:00Z",
    "text": "{\"cost\":3.1,\"tokensIn\":90000,\"tokensOut\":60000,\"cacheReads\":12000,\"cacheWrites\":0,\"apiProtocol\":\"anthropic\"}"
  }
]"#,
        )
        .unwrap();
        fs::write(
            architect_dir.join("api_conversation_history.json"),
            r#"before
<environment_details>
<model>claude-sonnet-4</model>
<slug>architect</slug>
<name>Architect</name>
</environment_details>
after"#,
        )
        .unwrap();

        let reviewer_dir = task_root.join("task-reviewer");
        fs::create_dir_all(&reviewer_dir).unwrap();
        fs::write(
            reviewer_dir.join("ui_messages.json"),
            r#"[
  {
    "type": "say",
    "say": "api_req_started",
    "ts": "2026-03-07T17:00:00Z",
    "text": "{\"cost\":1.8,\"tokensIn\":70000,\"tokensOut\":26000,\"cacheReads\":8000,\"cacheWrites\":0,\"apiProtocol\":\"anthropic\"}"
  },
  {
    "type": "say",
    "say": "api_req_started",
    "ts": "2026-03-07T17:09:00Z",
    "text": "{\"cost\":0.9,\"tokensIn\":22000,\"tokensOut\":18000,\"cacheReads\":3000,\"cacheWrites\":0,\"apiProtocol\":\"anthropic\"}"
  }
]"#,
        )
        .unwrap();
        fs::write(
            reviewer_dir.join("api_conversation_history.json"),
            r#"before
<environment_details>
<model>claude-haiku-4</model>
<slug>reviewer</slug>
<name>Reviewer</name>
</environment_details>
after"#,
        )
        .unwrap();

        unsafe {
            env::set_var("HOME", temp_dir.path());
        }

        let pricing = test_pricing_service();
        let loader = DataLoader::new(None);
        let usage = load_with_pricing(
            &loader,
            &[ClientId::RooCode],
            &GroupBy::Model,
            false,
            Some(&pricing),
        )
        .unwrap();

        let architect_expected = expected_message_cost(
            &pricing,
            "claude-sonnet-4",
            "anthropic",
            CoreTokenBreakdown {
                input: 420_000,
                output: 120_000,
                cache_read: 32_000,
                cache_write: 0,
                reasoning: 0,
            },
        ) + expected_message_cost(
            &pricing,
            "claude-sonnet-4",
            "anthropic",
            CoreTokenBreakdown {
                input: 90_000,
                output: 60_000,
                cache_read: 12_000,
                cache_write: 0,
                reasoning: 0,
            },
        );
        let reviewer_expected = expected_message_cost(
            &pricing,
            "claude-haiku-4",
            "anthropic",
            CoreTokenBreakdown {
                input: 70_000,
                output: 26_000,
                cache_read: 8_000,
                cache_write: 0,
                reasoning: 0,
            },
        ) + expected_message_cost(
            &pricing,
            "claude-haiku-4",
            "anthropic",
            CoreTokenBreakdown {
                input: 22_000,
                output: 18_000,
                cache_read: 3_000,
                cache_write: 0,
                reasoning: 0,
            },
        );

        assert_eq!(usage.agents.len(), 2);
        assert_eq!(usage.agents[0].agent, "Architect");
        assert_eq!(usage.agents[0].clients, "roocode");
        assert_eq!(usage.agents[0].message_count, 2);
        assert_cost_matches(usage.agents[0].cost, architect_expected);
        assert_eq!(usage.agents[0].tokens.total(), 734_000);

        assert_eq!(usage.agents[1].agent, "Reviewer");
        assert_eq!(usage.agents[1].message_count, 2);
        assert_cost_matches(usage.agents[1].cost, reviewer_expected);
        assert_eq!(usage.agents[1].tokens.total(), 147_000);

        match previous_home {
            Some(home) => unsafe { env::set_var("HOME", home) },
            None => unsafe { env::remove_var("HOME") },
        }
    }

    #[test]
    #[serial]
    fn test_data_loader_keeps_synthetic_gateway_messages_under_original_client() {
        let temp_dir = TempDir::new().unwrap();
        let previous_home = env::var_os("HOME");
        let previous_xdg_data_home = env::var_os("XDG_DATA_HOME");
        let message_dir = temp_dir
            .path()
            .join(".local/share/opencode/storage/message/project-1");
        fs::create_dir_all(&message_dir).unwrap();
        fs::write(
            message_dir.join("msg_001.json"),
            r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0.25,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#,
        )
        .unwrap();

        unsafe {
            env::set_var("HOME", temp_dir.path());
            env::remove_var("XDG_DATA_HOME");
        }

        let pricing = test_pricing_service();
        let loader = DataLoader::new(None);
        let usage = load_with_pricing(
            &loader,
            &[ClientId::OpenCode],
            &GroupBy::ClientProviderModel,
            true,
            Some(&pricing),
        )
        .unwrap();

        let expected_cost = expected_message_cost(
            &pricing,
            "accounts/fireworks/models/deepseek-v3-0324",
            "fireworks_ai",
            CoreTokenBreakdown {
                input: 10,
                output: 5,
                cache_read: 0,
                cache_write: 0,
                reasoning: 0,
            },
        );

        assert_eq!(usage.models.len(), 1);
        assert_eq!(usage.models[0].client, "opencode");
        // opencode now canonicalizes the provider (fireworks -> fireworks_ai),
        // matching every other session parser.
        assert_eq!(usage.models[0].provider, "fireworks_ai");
        assert_eq!(usage.models[0].model, "deepseek-v3-0324");
        assert_eq!(usage.models[0].tokens.total(), 15);
        assert_cost_matches(usage.models[0].cost, expected_cost);

        match previous_home {
            Some(home) => unsafe { env::set_var("HOME", home) },
            None => unsafe { env::remove_var("HOME") },
        }
        match previous_xdg_data_home {
            Some(path) => unsafe { env::set_var("XDG_DATA_HOME", path) },
            None => unsafe { env::remove_var("XDG_DATA_HOME") },
        }
    }

    #[test]
    fn test_calculate_streaks_uses_provided_today() {
        let today = NaiveDate::from_ymd_opt(2026, 3, 3).unwrap();
        let daily = vec![
            DailyUsage {
                date: NaiveDate::from_ymd_opt(2026, 3, 2).unwrap(),
                tokens: TokenBreakdown::default(),
                cost: 0.0,
                source_breakdown: BTreeMap::new(),
                message_count: 0,
                turn_count: 0,
            },
            DailyUsage {
                date: NaiveDate::from_ymd_opt(2026, 3, 3).unwrap(),
                tokens: TokenBreakdown::default(),
                cost: 0.0,
                source_breakdown: BTreeMap::new(),
                message_count: 0,
                turn_count: 0,
            },
        ];
        let (current, longest) = calculate_streaks_for_today(&daily, today);
        assert_eq!(current, 2);
        assert_eq!(longest, 2);
    }

    fn make_msg(timestamp_ms: i64, input: i64, output: i64, cost: f64) -> UnifiedMessage {
        UnifiedMessage::new(
            "claude",
            "claude-sonnet-4-5-20250929",
            "anthropic",
            format!("session-{timestamp_ms}"),
            timestamp_ms,
            tokmesh_core::TokenBreakdown {
                input,
                output,
                cache_read: 0,
                cache_write: 0,
                reasoning: 0,
            },
            cost,
        )
    }

    /// The Sessions tab's Model column reads `SessionUsage.models`. Each entry
    /// keeps the provider and color key its family-shade lookup needs, and the
    /// list is deduped by display name so a model reached under two spellings
    /// (here a dated id and its normalized form) lands once, in first-seen order.
    ///
    /// The dedup below relies on `model_name_for_grouping` ignoring `provider_id`
    /// for this input: it consults the provider only when `client == "opencode"`
    /// (to apply an OpenCode-configured label), and every other client falls
    /// through to `normalize_model_for_grouping`, which is a pure function of the
    /// model id. That is why the third message can carry a different provider
    /// (`github-copilot` vs `anthropic`) and still collapse onto the same display
    /// name. If that scoping ever widens beyond OpenCode, this assertion is the
    /// one that should start failing.
    #[test]
    fn test_session_models_dedup_by_display_name_and_retain_provider() {
        let loader = DataLoader::new(None);
        let base_ms = 1_735_689_600_000_i64;
        let tokens = || tokmesh_core::TokenBreakdown {
            input: 10,
            output: 5,
            cache_read: 0,
            cache_write: 0,
            reasoning: 0,
        };
        let usage = loader
            .aggregate_messages(
                vec![
                    UnifiedMessage::new(
                        "claude",
                        "claude-sonnet-4-5-20250929",
                        "anthropic",
                        "session-1",
                        base_ms,
                        tokens(),
                        1.0,
                    ),
                    UnifiedMessage::new(
                        "claude",
                        "gpt-5",
                        "openai",
                        "session-1",
                        base_ms + 1_000,
                        tokens(),
                        1.0,
                    ),
                    // Same model as the first message, spelled without the date
                    // and served through a gateway: must not add a second entry.
                    UnifiedMessage::new(
                        "claude",
                        "claude-sonnet-4-5",
                        "github-copilot",
                        "session-1",
                        base_ms + 2_000,
                        tokens(),
                        1.0,
                    ),
                ],
                &GroupBy::Model,
            )
            .unwrap();

        assert_eq!(usage.sessions.len(), 1, "expected a single session bucket");
        assert_eq!(
            usage.sessions[0].models,
            vec![
                SessionModel {
                    display_name: "claude-sonnet-4-5".to_string(),
                    provider: "anthropic".to_string(),
                    color_key: "claude-sonnet-4-5".to_string(),
                },
                SessionModel {
                    display_name: "gpt-5".to_string(),
                    provider: "openai".to_string(),
                    color_key: "gpt-5".to_string(),
                },
            ]
        );
    }

    #[test]
    fn test_minutely_aggregation_skipped_when_flag_disabled() {
        let loader = DataLoader::new(None);
        assert!(!loader.minutely_enabled);
        let usage = loader
            .aggregate_messages(
                vec![make_msg(1_735_689_600_000, 10, 5, 1.0)],
                &GroupBy::Model,
            )
            .unwrap();
        assert!(
            usage.minutely.is_empty(),
            "minutely aggregation should be skipped when flag is off"
        );
        assert_eq!(usage.hourly.len(), 1, "hourly should still aggregate");
    }

    #[test]
    fn test_minutely_aggregation_runs_when_flag_enabled() {
        let loader = DataLoader::new(None).with_minutely_enabled(true);
        assert!(loader.minutely_enabled);
        let usage = loader
            .aggregate_messages(
                vec![make_msg(1_735_689_600_000, 10, 5, 1.0)],
                &GroupBy::Model,
            )
            .unwrap();
        assert_eq!(usage.minutely.len(), 1);
        let bucket = &usage.minutely[0];
        assert_eq!(bucket.tokens.input, 10);
        assert_eq!(bucket.tokens.output, 5);
        assert_eq!(bucket.cost, 1.0);
        assert_eq!(bucket.message_count, 1);
    }

    #[test]
    fn test_minutely_aggregation_groups_same_minute_messages() {
        let loader = DataLoader::new(None).with_minutely_enabled(true);
        let base_ms = 1_735_689_600_000_i64;
        let usage = loader
            .aggregate_messages(
                vec![
                    make_msg(base_ms, 10, 5, 1.0),
                    make_msg(base_ms + 30_000, 20, 10, 2.0),
                ],
                &GroupBy::Model,
            )
            .unwrap();
        assert_eq!(
            usage.minutely.len(),
            1,
            "two messages within the same minute should share a bucket"
        );
        let bucket = &usage.minutely[0];
        assert_eq!(bucket.tokens.input, 30);
        assert_eq!(bucket.tokens.output, 15);
        assert_eq!(bucket.cost, 3.0);
        assert_eq!(bucket.message_count, 2);
    }

    #[test]
    fn test_minutely_aggregation_splits_adjacent_minutes() {
        let loader = DataLoader::new(None).with_minutely_enabled(true);
        let base_ms = 1_735_689_600_000_i64;
        let usage = loader
            .aggregate_messages(
                vec![
                    make_msg(base_ms, 10, 5, 1.0),
                    make_msg(base_ms + 60_000, 20, 10, 2.0),
                ],
                &GroupBy::Model,
            )
            .unwrap();
        assert_eq!(
            usage.minutely.len(),
            2,
            "messages one minute apart should land in distinct buckets"
        );
    }

    #[test]
    fn test_minutely_aggregation_clamps_negative_tokens_and_cost() {
        let loader = DataLoader::new(None).with_minutely_enabled(true);
        let usage = loader
            .aggregate_messages(
                vec![make_msg(1_735_689_600_000, -50, -50, -10.0)],
                &GroupBy::Model,
            )
            .unwrap();
        assert_eq!(usage.minutely.len(), 1);
        let bucket = &usage.minutely[0];
        assert_eq!(bucket.tokens.input, 0);
        assert_eq!(bucket.tokens.output, 0);
        assert_eq!(bucket.cost, 0.0);
    }

    fn daily_usage(
        date: NaiveDate,
        input: u64,
        output: u64,
        cost: f64,
        message_count: u32,
    ) -> DailyUsage {
        DailyUsage {
            date,
            tokens: TokenBreakdown {
                input,
                output,
                cache_read: 0,
                cache_write: 0,
                reasoning: 0,
            },
            cost,
            source_breakdown: BTreeMap::new(),
            message_count,
            turn_count: 0,
        }
    }

    #[test]
    fn test_aggregate_monthly_from_daily_groups_by_month() {
        let daily = vec![
            daily_usage(
                NaiveDate::from_ymd_opt(2026, 5, 10).unwrap(),
                100,
                50,
                1.0,
                2,
            ),
            daily_usage(
                NaiveDate::from_ymd_opt(2026, 5, 20).unwrap(),
                200,
                100,
                2.0,
                3,
            ),
            daily_usage(
                NaiveDate::from_ymd_opt(2026, 6, 1).unwrap(),
                300,
                150,
                3.0,
                4,
            ),
        ];

        let monthly = aggregate_monthly_from_daily(&daily);
        assert_eq!(monthly.len(), 2);

        let may = monthly.iter().find(|m| m.month == "2026-05").unwrap();
        assert_eq!(may.tokens.input, 300);
        assert_eq!(may.tokens.output, 150);
        assert_eq!(may.tokens.total(), 450);
        assert_eq!(may.cost, 3.0);
        assert_eq!(may.message_count, 5);

        let june = monthly.iter().find(|m| m.month == "2026-06").unwrap();
        assert_eq!(june.tokens.input, 300);
        assert_eq!(june.tokens.output, 150);
        assert_eq!(june.tokens.total(), 450);
        assert_eq!(june.cost, 3.0);
        assert_eq!(june.message_count, 4);
    }

    #[test]
    fn test_aggregate_messages_populates_monthly() {
        let loader = DataLoader::new(None);
        let base_ms = 1_736_899_200_000_i64; // mid-January 2025 in UTC
        let usage = loader
            .aggregate_messages(
                vec![
                    make_msg(base_ms, 100, 50, 1.0),
                    make_msg(base_ms + 86_400_000, 200, 100, 2.0), // +1 day
                ],
                &GroupBy::Model,
            )
            .unwrap();

        assert_eq!(usage.monthly.len(), 1);
        assert_eq!(usage.monthly[0].month, "2025-01");
        assert_eq!(usage.monthly[0].tokens.input, 300);
        assert_eq!(usage.monthly[0].cost, 3.0);
    }

    #[test]
    fn test_aggregate_monthly_sorts_descending() {
        let daily = vec![
            daily_usage(
                NaiveDate::from_ymd_opt(2026, 3, 1).unwrap(),
                100,
                50,
                1.0,
                1,
            ),
            daily_usage(
                NaiveDate::from_ymd_opt(2026, 5, 1).unwrap(),
                200,
                100,
                2.0,
                1,
            ),
        ];

        let monthly = aggregate_monthly_from_daily(&daily);
        assert_eq!(monthly[0].month, "2026-05");
        assert_eq!(monthly[1].month, "2026-03");
    }
}