claude-code-status-line 1.2.9

use crate::colors::SectionColors;
use crate::config::Config;
use crate::render::get_details_and_fg_codes;
use crate::types::{ContextAccuracy, ContextUsageInfo, ContextWindow};

pub fn calculate_context(
    context_window: Option<&ContextWindow>,
    config: &Config,
) -> Option<ContextUsageInfo> {
    let cw = context_window?;
    let total = cw.context_window_size;

    let display_remaining = config.sections.context.display_mode == "remaining";

    // Try to use direct percentages from Claude Code.
    let direct_percentage = if display_remaining {
        cw.remaining_percentage
    } else {
        cw.used_percentage
    };

    if let Some(pct) = direct_percentage {
        let used_tokens = match (cw.used_percentage, total) {
            (Some(used_pct), Some(total_tokens)) => {
                ((used_pct / 100.0) * total_tokens as f64) as u64
            }
            _ => 0,
        };

        return Some(ContextUsageInfo {
            percentage: pct.clamp(0.0, 100.0),
            current_usage_tokens: used_tokens,
            context_window_size: total,
            accuracy: ContextAccuracy::Exact,
        });
    }

    let usage = cw.current_usage.as_ref()?;
    let total = total?;

    if total == 0 {
        return None;
    }

    let used_tokens = usage.input_tokens.unwrap_or(0)
        + usage.cache_creation_input_tokens.unwrap_or(0)
        + usage.cache_read_input_tokens.unwrap_or(0);

    let used_pct = (used_tokens as f64 / total as f64 * 100.0).min(100.0);
    let percentage = if display_remaining {
        (100.0 - used_pct).max(0.0)
    } else {
        used_pct
    };

    Some(ContextUsageInfo {
        percentage,
        current_usage_tokens: used_tokens,
        context_window_size: Some(total),
        accuracy: ContextAccuracy::Estimated,
    })
}

pub fn format_ctx_short(info: &ContextUsageInfo, config: &Config) -> String {
    let prefix = match info.accuracy {
        ContextAccuracy::Exact => "",
        ContextAccuracy::Estimated => "~",
    };
    if config.sections.context.show_decimals {
        format!("ctx: {}{:.1}%", prefix, info.percentage)
    } else {
        format!("ctx: {}{}%", prefix, info.percentage as u64)
    }
}

pub fn format_ctx_full(info: &ContextUsageInfo, colors: &SectionColors, config: &Config) -> String {
    let (details, fg) = get_details_and_fg_codes(colors);
    let prefix = match info.accuracy {
        ContextAccuracy::Exact => "",
        ContextAccuracy::Estimated => "~",
    };

    let pct_str = if config.sections.context.show_decimals {
        format!("{}{:.1}%", prefix, info.percentage)
    } else {
        format!("{}{}%", prefix, info.percentage as u64)
    };

    if config.sections.context.show_token_counts
        && info.current_usage_tokens > 0
        && info.context_window_size.is_some()
    {
        let used_k = info.current_usage_tokens / 1000;
        let total_k = info.context_window_size.unwrap_or_default() / 1000;

        format!(
            "ctx: {} {}({}k/{}k){}",
            pct_str, details, used_k, total_k, fg
        )
    } else {
        format!("ctx: {}", pct_str)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ContextWindow, CurrentUsage};

    #[test]
    fn test_calculate_context_direct_used_percentage() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: None,
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: Some(45.5),
            remaining_percentage: Some(54.5),
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert_eq!(result.percentage, 45.5);
        assert_eq!(result.accuracy, ContextAccuracy::Exact);
        // 45.5% of 200k = 91k
        assert_eq!(result.current_usage_tokens, 91_000);
    }

    #[test]
    fn test_calculate_context_direct_remaining_percentage() {
        let mut config = Config::default();
        config.sections.context.display_mode = "remaining".to_string();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: None,
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: Some(45.5),
            remaining_percentage: Some(54.5),
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert_eq!(result.percentage, 54.5);
        assert_eq!(result.accuracy, ContextAccuracy::Exact);
    }

    #[test]
    fn test_calculate_context_fallback_used() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: Some(CurrentUsage {
                input_tokens: Some(50_000),
                output_tokens: Some(5_000),
                cache_creation_input_tokens: Some(10_000),
                cache_read_input_tokens: Some(5_000),
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert!(result.percentage > 0.0 && result.percentage <= 100.0);
        assert_eq!(result.accuracy, ContextAccuracy::Estimated);
        assert_eq!(result.context_window_size, Some(200_000));
        // 50k + 10k + 5k = 65k
        assert_eq!(result.current_usage_tokens, 65_000);
    }

    #[test]
    fn test_calculate_context_fallback_remaining() {
        let mut config = Config::default();
        config.sections.context.display_mode = "remaining".to_string();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: Some(CurrentUsage {
                input_tokens: Some(50_000),
                output_tokens: None,
                cache_creation_input_tokens: Some(10_000),
                cache_read_input_tokens: Some(5_000),
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        // 65k used = 32.5% used, so 67.5% remaining (allow for floating point precision)
        assert!((result.percentage - 67.5).abs() < 0.01);
        assert_eq!(result.accuracy, ContextAccuracy::Estimated);
    }

    #[test]
    fn test_calculate_context_zero_window() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(0),
            current_usage: Some(CurrentUsage {
                input_tokens: Some(100),
                output_tokens: None,
                cache_creation_input_tokens: None,
                cache_read_input_tokens: None,
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config);
        assert!(result.is_none());
    }

    #[test]
    fn test_calculate_context_missing_window() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: None,
            current_usage: Some(CurrentUsage {
                input_tokens: Some(50_000),
                output_tokens: None,
                cache_creation_input_tokens: None,
                cache_read_input_tokens: None,
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config);
        assert!(result.is_none());
    }

    #[test]
    fn test_calculate_context_exceeds_window() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(100_000),
            current_usage: Some(CurrentUsage {
                input_tokens: Some(95_000),
                output_tokens: None,
                cache_creation_input_tokens: Some(10_000),
                cache_read_input_tokens: None,
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert_eq!(result.percentage, 100.0);
    }

    #[test]
    fn test_calculate_context_all_zeros() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: Some(CurrentUsage {
                input_tokens: Some(0),
                output_tokens: Some(0),
                cache_creation_input_tokens: Some(0),
                cache_read_input_tokens: Some(0),
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert!(result.percentage < 25.0);
        assert_eq!(result.current_usage_tokens, 0);
    }

    #[test]
    fn test_calculate_context_missing_usage_no_direct() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: None,
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config);
        assert!(result.is_none());
    }

    #[test]
    fn test_calculate_context_null_tokens() {
        let config = Config::default();
        let cw = ContextWindow {
            context_window_size: Some(200_000),
            current_usage: Some(CurrentUsage {
                input_tokens: None,
                output_tokens: None,
                cache_creation_input_tokens: None,
                cache_read_input_tokens: None,
            }),
            total_input_tokens: None,
            total_output_tokens: None,
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert_eq!(result.current_usage_tokens, 0);
    }

    #[test]
    fn test_format_ctx_short_exact() {
        let config = Config::default();
        let result = format_ctx_short(
            &ContextUsageInfo {
                percentage: 42.7,
                current_usage_tokens: 0,
                context_window_size: Some(200_000),
                accuracy: ContextAccuracy::Exact,
            },
            &config,
        );
        assert_eq!(result, "ctx: 42%");
    }

    #[test]
    fn test_format_ctx_short_estimated() {
        let config = Config::default();
        let result = format_ctx_short(
            &ContextUsageInfo {
                percentage: 42.7,
                current_usage_tokens: 0,
                context_window_size: Some(200_000),
                accuracy: ContextAccuracy::Estimated,
            },
            &config,
        );
        assert_eq!(result, "ctx: ~42%");
    }

    #[test]
    fn test_format_ctx_short_exact_with_decimals() {
        let mut config = Config::default();
        config.sections.context.show_decimals = true;
        let result = format_ctx_short(
            &ContextUsageInfo {
                percentage: 42.7,
                current_usage_tokens: 0,
                context_window_size: Some(200_000),
                accuracy: ContextAccuracy::Exact,
            },
            &config,
        );
        assert_eq!(result, "ctx: 42.7%");
    }

    #[test]
    fn test_format_ctx_short_estimated_with_decimals() {
        let mut config = Config::default();
        config.sections.context.show_decimals = true;
        let result = format_ctx_short(
            &ContextUsageInfo {
                percentage: 42.7,
                current_usage_tokens: 0,
                context_window_size: Some(200_000),
                accuracy: ContextAccuracy::Estimated,
            },
            &config,
        );
        assert_eq!(result, "ctx: ~42.7%");
    }

    #[test]
    fn test_format_ctx_full_exact_no_token_counts() {
        let config = Config::default();
        let info = ContextUsageInfo {
            percentage: 55.5,
            current_usage_tokens: 0,
            context_window_size: Some(200_000),
            accuracy: ContextAccuracy::Exact,
        };

        let result = format_ctx_full(&info, &config.theme.context, &config);
        assert!(result.contains("ctx:"));
        assert!(result.contains("55%"));
        assert!(!result.contains("~"));
        assert!(!result.contains("k"));
    }

    #[test]
    fn test_format_ctx_full_exact_with_token_counts() {
        let mut config = Config::default();
        config.sections.context.show_token_counts = true;
        let info = ContextUsageInfo {
            percentage: 55.5,
            current_usage_tokens: 111_000,
            context_window_size: Some(200_000),
            accuracy: ContextAccuracy::Exact,
        };

        let result = format_ctx_full(&info, &config.theme.context, &config);
        assert!(result.contains("ctx:"));
        assert!(result.contains("55%"));
        assert!(!result.contains("~"));
        assert!(result.contains("111k"));
        assert!(result.contains("200k"));
    }

    #[test]
    fn test_format_ctx_full_estimated_with_token_counts() {
        let mut config = Config::default();
        config.sections.context.show_token_counts = true;
        let info = ContextUsageInfo {
            percentage: 55.5,
            current_usage_tokens: 111_000,
            context_window_size: Some(200_000),
            accuracy: ContextAccuracy::Estimated,
        };

        let result = format_ctx_full(&info, &config.theme.context, &config);
        assert!(result.contains("ctx:"));
        assert!(result.contains("~55%"));
        assert!(result.contains("111k"));
        assert!(result.contains("200k"));
    }

    #[test]
    fn test_format_ctx_full_estimated_without_token_counts() {
        let mut config = Config::default();
        config.sections.context.show_token_counts = false;
        let info = ContextUsageInfo {
            percentage: 55.5,
            current_usage_tokens: 111_000,
            context_window_size: Some(200_000),
            accuracy: ContextAccuracy::Estimated,
        };

        let result = format_ctx_full(&info, &config.theme.context, &config);
        assert!(result.contains("ctx:"));
        assert!(result.contains("~55%"));
        assert!(!result.contains("111k"));
        assert!(!result.contains("200k"));
    }

    #[test]
    fn test_format_ctx_full_direct_percentage_without_total() {
        let mut config = Config::default();
        config.sections.context.show_token_counts = true;
        let info = ContextUsageInfo {
            percentage: 12.0,
            current_usage_tokens: 0,
            context_window_size: None,
            accuracy: ContextAccuracy::Exact,
        };

        let result = format_ctx_full(&info, &config.theme.context, &config);
        assert!(result.contains("ctx: 12%"));
        assert!(!result.contains("k/"));
    }
}