magi-code 0.80.2

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()
}

/// One cooperative deadline shared by both sides of every file in a snapshot.
pub(crate) struct SourceHighlightBudget<'a> {
    pub(crate) deadline: std::time::Instant,
    pub(crate) cancellation: &'a crate::cancellation::AgentCancellation,
}

impl<'a> SourceHighlightBudget<'a> {
    pub(crate) fn new(cancellation: &'a crate::cancellation::AgentCancellation) -> Self {
        Self {
            deadline: std::time::Instant::now() + std::time::Duration::from_millis(500),
            cancellation,
        }
    }
}

#[cfg(test)]
pub(crate) fn highlight_source_file(input: &str, language: Option<&str>) -> Vec<DisplayLine> {
    highlight_source_with_budget(
        input,
        language,
        &SourceHighlightBudget::new(&crate::cancellation::AgentCancellation::default()),
    )
    .unwrap()
    .0
}

/// Worker-owned parser state. Deadlines pause parsing; they never discard progress.
pub(crate) struct SourceHighlightJob {
    source: String,
    lines: Vec<String>,
    display: std::sync::Arc<Vec<DisplayLine>>,
    parse_state: Option<ParseState>,
    scope_stack: ScopeStack,
    next_line: usize,
    initialized: bool,
    language: Option<String>,
}

impl SourceHighlightJob {
    pub(crate) fn new(source: String, language: Option<&str>) -> Self {
        let mut lines: Vec<String> = source.split('\n').map(str::to_owned).collect();
        if source.is_empty() || source.ends_with('\n') {
            lines.pop();
        }
        let display = std::sync::Arc::new(
            lines
                .iter()
                .map(|line| DisplayLine::from_span(line, DisplayRole::FallbackCode))
                .collect(),
        );
        Self {
            source,
            lines,
            display,
            parse_state: None,
            scope_stack: ScopeStack::new(),
            next_line: 0,
            initialized: false,
            language: language.map(str::to_owned),
        }
    }

    pub(crate) fn matches(&self, source: &str) -> bool {
        self.source == source
    }

    pub(crate) fn pending(&self) -> bool {
        self.next_line < self.lines.len()
    }

    pub(crate) fn display(&self) -> std::sync::Arc<Vec<DisplayLine>> {
        std::sync::Arc::clone(&self.display)
    }

    pub(crate) fn advance(
        &mut self,
        budget: &SourceHighlightBudget<'_>,
        max_lines: usize,
    ) -> anyhow::Result<()> {
        budget.cancellation.check()?;
        if !self.pending() || std::time::Instant::now() >= budget.deadline {
            return Ok(());
        }
        if !self.initialized {
            self.parse_state = self
                .language
                .as_deref()
                .filter(|_| self.source.len() <= 128 * 1024)
                .and_then(|language| {
                    find_syntax(&ASSETS.get_or_init(load_assets).syntaxes, language)
                })
                .map(ParseState::new);
            self.initialized = true;
        }
        let Some(parse_state) = self.parse_state.as_mut() else {
            self.next_line = self.lines.len();
            return Ok(());
        };
        let stop = self
            .next_line
            .saturating_add(max_lines)
            .min(self.lines.len());
        while self.next_line < stop && std::time::Instant::now() < budget.deadline {
            budget.cancellation.check()?;
            let line = &self.lines[self.next_line];
            // One parser call is not interruptible. Do not feed pathological long lines
            // or resume after skipping them: that would lose multiline syntax context.
            if line.len() > 4096 {
                self.next_line = self.lines.len();
                break;
            }
            let mut highlighted = highlight_line(
                &format!("{line}\n"),
                parse_state,
                &mut self.scope_stack,
                &ASSETS.get().expect("initialized syntax").syntaxes,
            );
            if let Some(span) = highlighted.spans.last_mut()
                && span.text.ends_with('\n')
            {
                span.text.pop();
            }
            std::sync::Arc::make_mut(&mut self.display)[self.next_line] = highlighted;
            self.next_line += 1;
            budget.cancellation.check()?;
        }
        Ok(())
    }
}

#[cfg(test)]
pub(crate) fn highlight_source_with_budget(
    input: &str,
    language: Option<&str>,
    budget: &SourceHighlightBudget<'_>,
) -> anyhow::Result<(Vec<DisplayLine>, bool)> {
    let mut job = SourceHighlightJob::new(input.to_owned(), language);
    job.advance(budget, usize::MAX)?;
    let mut lines = (*job.display()).clone();
    if input.is_empty() || input.ends_with('\n') {
        lines.push(DisplayLine::from_span("", DisplayRole::FallbackCode));
    }
    Ok((lines, job.pending()))
}

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, table: None }
    }
}

fn semantic_scope_rules() -> &'static [SemanticScopeRule] {
    SEMANTIC_SCOPE_RULES
        .get_or_init(|| {
            [
                (DisplayRole::Comment, "comment, punctuation.definition.comment"),
                (DisplayRole::String, "string, constant.character, punctuation.definition.string"),
                (DisplayRole::Number, "constant.numeric"),
                (DisplayRole::Keyword, "keyword, storage.modifier, 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::Macro, "support.macro, entity.name.macro"),
                (DisplayRole::Attribute, "meta.annotation, variable.annotation"),
                (DisplayRole::Lifetime, "storage.modifier.lifetime"),
                (DisplayRole::Field, "variable.other.member"),
                (DisplayRole::Operator, "keyword.operator"),
                (DisplayRole::Punctuation, "punctuation"),
                (DisplayRole::Heading, "markup.heading, entity.name.section, punctuation.definition.heading"),
                (DisplayRole::Strong, "markup.bold, punctuation.definition.bold"),
                (DisplayRole::Emphasis, "markup.italic, punctuation.definition.italic"),
                (DisplayRole::InlineCode, "markup.raw, punctuation.definition.raw"),
                (DisplayRole::CodeFence, "punctuation.definition.raw.code-fence"),
                (DisplayRole::CodeLanguageLabel, "constant.other.language-name"),
                (DisplayRole::Link, "meta.link, markup.underline.link, punctuation.definition.link"),
                (DisplayRole::ListMarker, "punctuation.definition.list_item"),
                (DisplayRole::BlockQuote, "markup.quote, punctuation.definition.blockquote"),
            ]
            .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::*;
    fn assert_token_role(source: &str, language: &str, token: &str, role: DisplayRole) {
        let lines = highlight_source_file(source, Some(language));
        assert_eq!(super::super::plain_projection(&lines), source);
        assert!(
            lines
                .iter()
                .flat_map(|line| &line.spans)
                .any(|span| span.text.contains(token) && span.role == role),
            "{token:?} should be {role:?}: {lines:?}"
        );
    }

    #[test]
    fn rust_source_maps_scoped_tokens() {
        let source = "/// Documentation\n#[derive(Debug)]\nstruct Widget<'a> { field: &'a str }\nfn main() { println!(\"hello\"); item.run(); }";
        for (token, role) in [
            ("///", DisplayRole::Comment),
            ("Documentation", DisplayRole::Comment),
            ("derive", DisplayRole::Attribute),
            ("Debug", DisplayRole::Attribute),
            ("Widget", DisplayRole::Type),
            ("'a", DisplayRole::Lifetime),
            ("field", DisplayRole::Field),
            ("main", DisplayRole::Function),
            ("println!", DisplayRole::Macro),
            ("run", DisplayRole::Function),
            ("{", DisplayRole::Punctuation),
            ("hello", DisplayRole::String),
        ] {
            assert_token_role(source, "rs", token, role);
        }
    }

    #[test]
    fn markdown_source_maps_structure_and_inline_tokens() {
        let source = "# Heading\n\n**bold** *italic* `code` [label](https://example.org)\n\n- list\n\n> quote\n\n```rust\nfn main() {}\n```";
        for (token, role) in [
            ("Heading", DisplayRole::Heading),
            ("bold", DisplayRole::Strong),
            ("italic", DisplayRole::Emphasis),
            ("code", DisplayRole::InlineCode),
            ("label", DisplayRole::Link),
            ("https://example.org", DisplayRole::Link),
            ("-", DisplayRole::ListMarker),
            ("quote", DisplayRole::BlockQuote),
            ("```", DisplayRole::CodeFence),
            ("rust", DisplayRole::CodeLanguageLabel),
        ] {
            assert_token_role(source, "md", token, role);
        }
    }

    #[test]
    fn source_highlighting_preserves_multiline_state_beyond_transcript_limit() {
        let source = format!(
            "/*\n{}*/\nfn after_comment() {{}}",
            "documentation\n".repeat(450)
        );
        let started = std::time::Instant::now();
        let lines = highlight_source_file(&source, Some("rs"));
        eprintln!(
            "source fixture: {} bytes, {} lines, {:?}",
            source.len(),
            lines.len(),
            started.elapsed()
        );
        assert_eq!(super::super::plain_projection(&lines), source);
        assert!(
            lines[440]
                .spans
                .iter()
                .any(|span| span.role == DisplayRole::Comment)
        );
        assert!(
            lines
                .last()
                .unwrap()
                .spans
                .iter()
                .any(|span| span.text == "after_comment" && span.role == DisplayRole::Function)
        );
    }

    #[test]
    fn source_budget_fallback_preserves_remaining_text() {
        let source = format!(
            "fn highlighted() {{}}\n{}\nfn fallback() {{}}",
            "x".repeat(4097)
        );
        let lines = highlight_source_file(&source, Some("rs"));
        assert_eq!(super::super::plain_projection(&lines), source);
        assert!(
            lines[0]
                .spans
                .iter()
                .any(|span| span.role == DisplayRole::Function)
        );
        assert!(
            lines[1..]
                .iter()
                .flat_map(|line| &line.spans)
                .all(|span| span.role == DisplayRole::FallbackCode)
        );
    }

    #[test]
    #[ignore = "reports worker highlighting cost on repository source fixtures"]
    fn profile_source_highlighting() {
        for pass in ["cold", "warm"] {
            for (language, source) in [
                ("rs", include_str!("../tui/theme.rs")),
                (
                    "md",
                    include_str!("../../docs/features/mission-control-tui.md"),
                ),
            ] {
                let started = std::time::Instant::now();
                let lines = highlight_source_file(source, Some(language));
                let colored = lines
                    .iter()
                    .filter(|line| {
                        line.spans
                            .iter()
                            .any(|span| span.role != DisplayRole::FallbackCode)
                    })
                    .count();
                eprintln!(
                    "{pass} {language}: {} bytes, {colored}/{} lines with roles, {:?}",
                    source.len(),
                    lines.len(),
                    started.elapsed()
                );
                assert_eq!(super::super::plain_projection(&lines), source);
            }
        }
    }

    #[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 resumable_source_completes_beyond_former_line_cutoff() {
        prewarm();
        let source = format!("{}fn final_line() {{}}", "// x\n".repeat(8200));
        let token = crate::cancellation::AgentCancellation::default();
        let mut job = SourceHighlightJob::new(source, Some("rs"));
        let started = std::time::Instant::now();
        let mut passes = 0;
        while job.pending() {
            job.advance(&SourceHighlightBudget::new(&token), 1000)
                .unwrap();
            passes += 1;
            assert!(passes < 100);
        }
        assert!(
            job.display()[8200]
                .spans
                .iter()
                .any(|span| span.text == "final_line" && span.role == DisplayRole::Function)
        );
        eprintln!(
            "8201 lines completed in {passes} slices: {:?}",
            started.elapsed()
        );
    }

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

#[cfg(test)]
mod cancellation_tests {
    use super::*;
    use crate::cancellation::AgentCancellation;
    use std::time::{Duration, Instant};

    #[test]
    fn cancellation_during_source_projection_stops_before_deadline() {
        prewarm();
        let (token, cancel) = AgentCancellation::default().child_token();
        let source = "fn main() { let number = 123; }\n".repeat(100_000);
        let worker = std::thread::spawn(move || {
            highlight_source_with_budget(
                &source,
                Some("rs"),
                &SourceHighlightBudget {
                    deadline: Instant::now() + Duration::from_secs(30),
                    cancellation: &token,
                },
            )
        });
        std::thread::sleep(Duration::from_millis(5));
        cancel.cancel();
        assert!(worker.join().unwrap().is_err());
    }
}