magi-code 0.63.1

use super::{DisplayLine, DisplayRole, DisplaySpan};
use std::sync::OnceLock;
use syntect::{
    easy::HighlightLines,
    highlighting::{Style as SyntectStyle, Theme, ThemeSet},
    parsing::{SyntaxReference, SyntaxSet},
};

const MAX_HIGHLIGHT_LINES: usize = 400;
const MAX_HIGHLIGHT_BYTES: usize = 64 * 1024;

struct HighlightAssets {
    syntaxes: SyntaxSet,
    theme: Theme,
}

static ASSETS: OnceLock<HighlightAssets> = OnceLock::new();

/// Eagerly initialize syntect highlighting assets to avoid cold-start latency
/// on the first highlighted code block during TUI rendering.
pub(crate) fn prewarm() {
    ASSETS.get_or_init(load_assets);
}

pub(crate) fn highlight_code(input: &str, language: Option<&str>) -> Vec<DisplayLine> {
    if input.is_empty() {
        return vec![DisplayLine::from_span("", DisplayRole::FallbackCode)];
    }
    let line_count = input.split('\n').count();
    if line_count > MAX_HIGHLIGHT_LINES || input.len() > MAX_HIGHLIGHT_BYTES {
        return fallback_code(input);
    }

    let Some(language) = language.filter(|language| !language.trim().is_empty()) else {
        return fallback_code(input);
    };
    let assets = ASSETS.get_or_init(load_assets);
    let Some(syntax) = find_syntax(&assets.syntaxes, language) else {
        return fallback_code(input);
    };

    let mut highlighter = HighlightLines::new(syntax, &assets.theme);
    input
        .split('\n')
        .map(|line| highlight_line(line, &mut highlighter, &assets.syntaxes))
        .collect()
}

fn load_assets() -> HighlightAssets {
    let syntaxes = SyntaxSet::load_defaults_newlines();
    let themes = ThemeSet::load_defaults();
    let theme = themes
        .themes
        .get("InspiredGitHub")
        .or_else(|| themes.themes.values().next())
        .cloned()
        .unwrap_or_default();
    HighlightAssets { syntaxes, theme }
}

fn find_syntax<'a>(syntaxes: &'a SyntaxSet, language: &str) -> Option<&'a SyntaxReference> {
    let language = language.trim().trim_start_matches('.');
    syntaxes
        .find_syntax_by_token(language)
        .or_else(|| syntaxes.find_syntax_by_extension(language))
        .or_else(|| syntaxes.find_syntax_by_name(language))
}

fn highlight_line(
    line: &str,
    highlighter: &mut HighlightLines<'_>,
    syntaxes: &SyntaxSet,
) -> DisplayLine {
    let highlighted = highlighter.highlight_line(line, syntaxes);
    match highlighted {
        Ok(ranges) => {
            let spans = ranges
                .into_iter()
                .filter(|(_, text)| !text.is_empty())
                .map(|(style, text)| DisplaySpan::new(text, role_for_style(style)))
                .collect::<Vec<_>>();
            if spans.is_empty() {
                DisplayLine::from_span("", DisplayRole::FallbackCode)
            } else {
                DisplayLine { spans }
            }
        }
        Err(_) => DisplayLine::from_span(line, DisplayRole::FallbackCode),
    }
}

/// Maps syntect `Style` foreground RGB values to display color buckets.
///
/// This is not syntax-scope mapping: `HighlightLines::highlight_line` returns theme
/// colors, not the token `ScopeStack`. Buckets are calibrated for the
/// `InspiredGitHub` theme selected in `load_assets()`, so role names here are
/// rendering hints and may not match Rust token semantics.
fn role_for_style(style: SyntectStyle) -> DisplayRole {
    let fg = style.foreground;
    if fg.r > 150 && fg.g < 120 && fg.b < 120 {
        DisplayRole::String
    } else if fg.r < 120 && fg.g < 140 && fg.b < 160 {
        DisplayRole::Comment
    } else if fg.b > 150 && fg.r < 120 {
        DisplayRole::Keyword
    } else if fg.b > 150 && fg.r < 170 {
        DisplayRole::Function
    } else if fg.r > 130 && fg.b > 130 && fg.g < 150 {
        DisplayRole::Type
    } else if fg.r > 130 && fg.g > 100 && fg.b < 120 {
        DisplayRole::Number
    } else if fg.r == fg.g && fg.g == fg.b {
        DisplayRole::Punctuation
    } else if fg.g > fg.r && fg.g > fg.b {
        DisplayRole::Operator
    } else {
        DisplayRole::FallbackCode
    }
}

fn fallback_code(input: &str) -> Vec<DisplayLine> {
    input
        .split('\n')
        .map(|line| DisplayLine::from_span(line, DisplayRole::FallbackCode))
        .collect()
}

#[cfg(test)]
fn roles_for(input: &str, language: Option<&str>) -> Vec<DisplayRole> {
    highlight_code(input, language)
        .into_iter()
        .flat_map(|line| line.spans)
        .map(|span| span.role)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn prewarm_initializes_assets() {
        prewarm();
        assert!(
            ASSETS.get().is_some(),
            "ASSETS must be initialized after prewarm()"
        );
    }

    #[test]
    fn rust_keyword_uses_a_non_fallback_color_bucket() {
        let roles = roles_for("fn", Some("rust"));
        assert!(
            roles.iter().any(|role| *role != DisplayRole::FallbackCode),
            "recognized Rust keyword should receive some theme color bucket"
        );
    }

    #[test]
    fn rust_string_literal_uses_a_non_fallback_color_bucket() {
        let roles = roles_for("\"hello\"", Some("rust"));
        assert!(
            roles.iter().any(|role| *role != DisplayRole::FallbackCode),
            "recognized Rust string literal should receive some theme color bucket"
        );
    }

    #[test]
    fn rust_comment_uses_a_non_fallback_color_bucket() {
        let roles = roles_for("// comment", Some("rust"));
        assert!(
            roles.iter().any(|role| *role != DisplayRole::FallbackCode),
            "recognized Rust comment should receive some theme color bucket"
        );
    }

    #[test]
    fn highlight_role_mapping_is_theme_color_bucket_not_semantic_scope() {
        assert!(
            roles_for("fn", Some("rust")).contains(&DisplayRole::String),
            "InspiredGitHub colors currently bucket `fn` as String; this is not a semantic keyword assertion"
        );
        assert!(
            roles_for("\"hello\"", Some("rust")).contains(&DisplayRole::Comment),
            "InspiredGitHub colors currently bucket string literals as Comment; this is not a semantic string assertion"
        );
        assert!(
            roles_for("// comment", Some("rust")).contains(&DisplayRole::Operator),
            "InspiredGitHub colors currently bucket comments as Operator; this is not a semantic comment assertion"
        );
    }

    #[test]
    fn recognized_language_uses_non_fallback_roles() {
        let lines = highlight_code("fn main() {\n    let n = 1;\n}", Some("rust"));
        assert!(
            lines
                .iter()
                .flat_map(|line| &line.spans)
                .any(|span| span.role != DisplayRole::FallbackCode)
        );
    }

    #[test]
    fn unknown_and_missing_language_fall_back() {
        for language in [Some("not-a-real-language"), None] {
            let lines = highlight_code("let x = 1;", language);
            assert!(
                lines
                    .iter()
                    .flat_map(|line| &line.spans)
                    .all(|span| span.role == DisplayRole::FallbackCode)
            );
        }
    }

    #[test]
    fn budget_overflow_falls_back_without_panic() {
        let input = (0..=MAX_HIGHLIGHT_LINES)
            .map(|_| "fn main() {}")
            .collect::<Vec<_>>()
            .join("\n");
        let lines = highlight_code(&input, Some("rust"));
        assert!(
            lines
                .iter()
                .flat_map(|line| &line.spans)
                .all(|span| span.role == DisplayRole::FallbackCode)
        );
    }
}