claude-code-status-line 1.2.5

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

/// Fallback context window size (only used if Claude Code doesn't provide it)
const DEFAULT_CONTEXT_WINDOW: u64 = 200_000;

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

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

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

    if let Some(pct) = direct_percentage {
        // We have exact percentage from Claude Code
        return Some(ContextUsageInfo {
            percentage: pct.clamp(0.0, 100.0),
            current_usage_tokens: 0, // Not available when using direct percentage
            context_window_size: total,
            is_exact: true,
        });
    }

    // Fall back to calculated percentage for older Claude Code versions
    if total == 0 {
        return None;
    }
    let usage = cw.current_usage.as_ref()?;

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

    let total_used = content_tokens;
    let used_pct = (total_used 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: total_used,
        context_window_size: total,
        is_exact: false,
    })
}

pub fn format_ctx_short(pct: f64, is_exact: bool, config: &Config) -> String {
    let prefix = if is_exact { "" } else { "~" };
    if config.sections.context.show_decimals {
        format!("ctx: {}{:.1}%", prefix, pct)
    } else {
        format!("ctx: {}{}%", prefix, pct as u64)
    }
}

pub fn format_ctx_full(info: &ContextUsageInfo, colors: &SectionColors, config: &Config) -> String {
    let (details, fg) = get_details_and_fg_codes(colors, config);
    let prefix = if info.is_exact { "" } else { "~" };

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

    // Only show token counts if we have them (fallback mode) and config enables it
    if config.sections.context.show_token_counts && !info.is_exact && info.current_usage_tokens > 0
    {
        let used_k = info.current_usage_tokens / 1000;
        let total_k = info.context_window_size / 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,
            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!(result.is_exact);
        assert_eq!(result.current_usage_tokens, 0);
    }

    #[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,
            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!(result.is_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),
            }),
            used_percentage: None,
            remaining_percentage: None,
        };

        let result = calculate_context(Some(&cw), &config).unwrap();
        assert!(result.percentage > 0.0 && result.percentage <= 100.0);
        assert!(!result.is_exact);
        assert_eq!(result.context_window_size, 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),
            }),
            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!(!result.is_exact);
    }

    #[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,
            }),
            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,
            }),
            used_percentage: None,
            remaining_percentage: None,
        };

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

    #[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,
            }),
            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),
            }),
            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,
            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,
            }),
            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(42.7, true, &config);
        assert_eq!(result, "ctx: 42%");
    }

    #[test]
    fn test_format_ctx_short_estimated() {
        let config = Config::default();
        let result = format_ctx_short(42.7, false, &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(42.7, true, &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(42.7, false, &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: 200_000,
            is_exact: true,
        };

        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_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: 200_000,
            is_exact: false,
        };

        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: 200_000,
            is_exact: false,
        };

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