magi-code 0.77.1

use super::{DisplayLine, DisplayRole, DisplaySpan};
use std::{str::FromStr, sync::OnceLock};
use syntect::{
    easy::ScopeRegionIterator,
    highlighting::ScopeSelectors,
    parsing::{ParseState, ScopeStack, SyntaxReference, SyntaxSet},
};

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

struct HighlightAssets {
    syntaxes: SyntaxSet,
}

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

#[derive(Debug)]
struct SemanticScopeRule {
    role: DisplayRole,
    selectors: ScopeSelectors,
}

static SEMANTIC_SCOPE_RULES: OnceLock<Vec<SemanticScopeRule>> = OnceLock::new();

/// Eagerly initialize syntect syntax assets to avoid cold-start latency on the
/// first highlighted code block during TUI rendering.
#[cfg(test)]
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 parse_state = ParseState::new(syntax);
    let mut scope_stack = ScopeStack::new();
    input
        .split('\n')
        .map(|line| highlight_line(line, &mut parse_state, &mut scope_stack, &assets.syntaxes))
        .collect()
}

fn load_assets() -> HighlightAssets {
    HighlightAssets {
        syntaxes: SyntaxSet::load_defaults_newlines(),
    }
}

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,
    parse_state: &mut ParseState,
    scope_stack: &mut ScopeStack,
    syntaxes: &SyntaxSet,
) -> DisplayLine {
    let Ok(operations) = parse_state.parse_line(line, syntaxes) else {
        return DisplayLine::from_span(line, DisplayRole::FallbackCode);
    };

    let mut spans = Vec::new();
    for (text, operation) in ScopeRegionIterator::new(&operations, line) {
        if scope_stack.apply(operation).is_err() {
            return DisplayLine::from_span(line, DisplayRole::FallbackCode);
        }
        if text.is_empty() {
            continue;
        }
        spans.push(DisplaySpan::new(
            text,
            role_for_scopes(scope_stack.as_slice()),
        ));
    }

    if spans.is_empty() {
        DisplayLine::from_span("", DisplayRole::FallbackCode)
    } else {
        DisplayLine { spans }
    }
}

fn semantic_scope_rules() -> &'static [SemanticScopeRule] {
    SEMANTIC_SCOPE_RULES
        .get_or_init(|| {
            [
                (DisplayRole::Comment, "comment"),
                (DisplayRole::String, "string, constant.character"),
                (DisplayRole::Number, "constant.numeric"),
                (DisplayRole::Keyword, "keyword, storage.type.function"),
                (
                    DisplayRole::Function,
                    "entity.name.function, support.function",
                ),
                (
                    DisplayRole::Type,
                    "entity.name.type, entity.name.class, entity.name.struct, entity.name.enum, storage.type, support.type",
                ),
                (DisplayRole::Operator, "keyword.operator"),
                (DisplayRole::Punctuation, "punctuation"),
            ]
            .into_iter()
            .map(|(role, selector)| SemanticScopeRule {
                role,
                selectors: ScopeSelectors::from_str(selector)
                    .expect("built-in semantic syntax selector must be valid"),
            })
            .collect()
        })
        .as_slice()
}

/// Maps syntax-definition scopes to app-owned semantic display roles.
///
/// Syntax definitions identify what a token is; the active TUI theme decides
/// how that role is rendered. Keeping this mapping independent of syntect's
/// optional color themes prevents a fixed RGB palette from changing semantics.
fn role_for_scopes(scopes: &[syntect::parsing::Scope]) -> DisplayRole {
    semantic_scope_rules()
        .iter()
        .filter_map(|rule| {
            rule.selectors
                .does_match(scopes)
                .map(|power| (power, rule.role))
        })
        .max_by_key(|(power, _)| *power)
        .map(|(_, role)| role)
        .unwrap_or(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_maps_to_keyword_role() {
        assert!(
            roles_for("fn", Some("rust")).contains(&DisplayRole::Keyword),
            "recognized Rust keyword should use the semantic keyword role"
        );
    }

    #[test]
    fn rust_string_literal_maps_to_string_role() {
        assert!(
            roles_for("\"hello\"", Some("rust")).contains(&DisplayRole::String),
            "recognized Rust string literal should use the semantic string role"
        );
    }

    #[test]
    fn rust_comment_maps_to_comment_role() {
        assert!(
            roles_for("// comment", Some("rust")).contains(&DisplayRole::Comment),
            "recognized Rust comment should use the semantic comment role"
        );
    }

    #[test]
    fn rust_function_name_maps_to_function_role() {
        assert!(roles_for("fn main", Some("rust")).contains(&DisplayRole::Function));
    }

    #[test]
    fn rust_type_name_maps_to_type_role() {
        assert!(roles_for("struct Widget", Some("rust")).contains(&DisplayRole::Type));
    }

    #[test]
    fn rust_number_maps_to_number_role() {
        assert!(roles_for("let count = 42;", Some("rust")).contains(&DisplayRole::Number));
    }

    #[test]
    fn rust_operator_maps_to_operator_role() {
        assert!(roles_for("left + right", Some("rust")).contains(&DisplayRole::Operator));
    }

    #[test]
    fn rust_punctuation_maps_to_punctuation_role() {
        assert!(roles_for("let value = 1;", Some("rust")).contains(&DisplayRole::Punctuation));
    }

    #[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)
        );
    }
}