Cagire/src/views/highlight.rs

use ratatui::style::{Modifier, Style};

use crate::model::{SourceSpan, WordCompile, WORDS};
use crate::theme::syntax;

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum TokenKind {
    Number,
    String,
    Comment,
    Keyword,
    StackOp,
    Operator,
    Sound,
    Param,
    Context,
    Note,
    Interval,
    Variable,
    Emit,
    Vary,
    Generator,
    Default,
}

impl TokenKind {
    pub fn style(self) -> Style {
        let (fg, bg) = match self {
            TokenKind::Emit => syntax::EMIT,
            TokenKind::Number => syntax::NUMBER,
            TokenKind::String => syntax::STRING,
            TokenKind::Comment => syntax::COMMENT,
            TokenKind::Keyword => syntax::KEYWORD,
            TokenKind::StackOp => syntax::STACK_OP,
            TokenKind::Operator => syntax::OPERATOR,
            TokenKind::Sound => syntax::SOUND,
            TokenKind::Param => syntax::PARAM,
            TokenKind::Context => syntax::CONTEXT,
            TokenKind::Note => syntax::NOTE,
            TokenKind::Interval => syntax::INTERVAL,
            TokenKind::Variable => syntax::VARIABLE,
            TokenKind::Vary => syntax::VARY,
            TokenKind::Generator => syntax::GENERATOR,
            TokenKind::Default => syntax::DEFAULT,
        };
        let style = Style::default().fg(fg).bg(bg);
        if matches!(self, TokenKind::Emit) {
            style.add_modifier(Modifier::BOLD)
        } else {
            style
        }
    }

    pub fn gap_style() -> Style {
        Style::default().bg(syntax::GAP_BG)
    }
}

pub struct Token {
    pub start: usize,
    pub end: usize,
    pub kind: TokenKind,
    pub varargs: bool,
}

fn lookup_word_kind(word: &str) -> Option<(TokenKind, bool)> {
    if word == "." {
        return Some((TokenKind::Emit, false));
    }
    if word == ".!" {
        return Some((TokenKind::Emit, true));
    }

    for w in WORDS {
        if w.name == word || w.aliases.contains(&word) {
            let kind = match &w.compile {
                WordCompile::Param => TokenKind::Param,
                WordCompile::Context(_) => TokenKind::Context,
                _ => match w.category {
                    "Stack" => TokenKind::StackOp,
                    "Arithmetic" | "Comparison" | "Music" => TokenKind::Operator,
                    "Logic" if matches!(w.name, "and" | "or" | "not" | "xor" | "nand" | "nor") => {
                        TokenKind::Operator
                    }
                    "Sound" => TokenKind::Sound,
                    "Randomness" | "Probability" | "Selection" => TokenKind::Vary,
                    "Generator" => TokenKind::Generator,
                    _ => TokenKind::Keyword,
                },
            };
            return Some((kind, w.varargs));
        }
    }
    None
}

fn is_note(word: &str) -> bool {
    let bytes = word.as_bytes();
    if bytes.len() < 2 {
        return false;
    }
    let base = matches!(bytes[0], b'c' | b'd' | b'e' | b'f' | b'g' | b'a' | b'b');
    if !base {
        return false;
    }
    match bytes[1] {
        b'#' | b's' | b'b' => bytes.len() > 2 && bytes[2..].iter().all(|b| b.is_ascii_digit()),
        b'0'..=b'9' => bytes[1..].iter().all(|b| b.is_ascii_digit()),
        _ => false,
    }
}

const INTERVALS: &[&str] = &[
    "P1", "unison", "m2", "M2", "m3", "M3", "P4", "aug4", "dim5", "tritone", "P5", "m6", "M6",
    "m7", "M7", "P8", "oct", "m9", "M9", "m10", "M10", "P11", "aug11", "P12", "m13", "M13", "m14",
    "M14", "P15",
];

pub fn tokenize_line(line: &str) -> Vec<Token> {
    let mut tokens = Vec::new();
    let mut chars = line.char_indices().peekable();

    while let Some((start, c)) = chars.next() {
        if c.is_whitespace() {
            continue;
        }

        if c == ';' && chars.peek().map(|(_, ch)| *ch) == Some(';') {
            tokens.push(Token {
                start,
                end: line.len(),
                kind: TokenKind::Comment,
                varargs: false,
            });
            break;
        }

        if c == '"' {
            let mut end = start + 1;
            for (i, ch) in chars.by_ref() {
                end = i + ch.len_utf8();
                if ch == '"' {
                    break;
                }
            }
            tokens.push(Token {
                start,
                end,
                kind: TokenKind::String,
                varargs: false,
            });
            continue;
        }

        let mut end = start + c.len_utf8();
        while let Some((i, ch)) = chars.peek() {
            if ch.is_whitespace() {
                break;
            }
            end = *i + ch.len_utf8();
            chars.next();
        }

        let word = &line[start..end];
        let (kind, varargs) = classify_word(word);
        tokens.push(Token { start, end, kind, varargs });
    }

    tokens
}

fn classify_word(word: &str) -> (TokenKind, bool) {
    if word.parse::<f64>().is_ok() || word.parse::<i64>().is_ok() {
        return (TokenKind::Number, false);
    }

    if let Some((kind, varargs)) = lookup_word_kind(word) {
        return (kind, varargs);
    }

    if INTERVALS.contains(&word) {
        return (TokenKind::Interval, false);
    }

    if is_note(&word.to_ascii_lowercase()) {
        return (TokenKind::Note, false);
    }

    if word.len() > 1 && (word.starts_with('@') || word.starts_with('!')) {
        return (TokenKind::Variable, false);
    }

    (TokenKind::Default, false)
}

pub fn highlight_line(line: &str) -> Vec<(Style, String)> {
    highlight_line_with_runtime(line, &[], &[])
}

pub fn highlight_line_with_runtime(
    line: &str,
    executed_spans: &[SourceSpan],
    selected_spans: &[SourceSpan],
) -> Vec<(Style, String)> {
    let tokens = tokenize_line(line);
    let mut result = Vec::new();
    let mut last_end = 0;
    let gap_style = TokenKind::gap_style();

    for token in tokens {
        if token.start > last_end {
            result.push((gap_style, line[last_end..token.start].to_string()));
        }

        let is_selected = selected_spans
            .iter()
            .any(|span| overlaps(token.start, token.end, span.start, span.end));
        let is_executed = executed_spans
            .iter()
            .any(|span| overlaps(token.start, token.end, span.start, span.end));

        let mut style = token.kind.style();
        if token.varargs {
            style = style.add_modifier(Modifier::UNDERLINED);
        }
        if is_selected {
            style = style.bg(syntax::SELECTED_BG).add_modifier(Modifier::BOLD);
        } else if is_executed {
            style = style.bg(syntax::EXECUTED_BG);
        }

        result.push((style, line[token.start..token.end].to_string()));
        last_end = token.end;
    }

    if last_end < line.len() {
        result.push((gap_style, line[last_end..].to_string()));
    }

    result
}

fn overlaps(a_start: usize, a_end: usize, b_start: usize, b_end: usize) -> bool {
    a_start < b_end && b_start < a_end
}