magi-code 0.96.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();

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)]
mod profiling {
    use super::*;
    #[test]
    #[ignore = "reports worker highlighting cost on repository source fixtures; run with --release --ignored --nocapture"]
    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);
            }
        }
    }
}