Cagire/src/views/highlight.rs

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

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

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

impl TokenKind {
    pub fn style(self) -> Style {
        match self {
            TokenKind::Number => Style::default().fg(Color::Rgb(255, 180, 100)),
            TokenKind::String => Style::default().fg(Color::Rgb(150, 220, 150)),
            TokenKind::Comment => Style::default().fg(Color::Rgb(100, 100, 100)),
            TokenKind::Keyword => Style::default().fg(Color::Rgb(220, 120, 220)),
            TokenKind::StackOp => Style::default().fg(Color::Rgb(120, 180, 220)),
            TokenKind::Operator => Style::default().fg(Color::Rgb(200, 200, 130)),
            TokenKind::Sound => Style::default().fg(Color::Rgb(100, 220, 200)),
            TokenKind::Param => Style::default().fg(Color::Rgb(180, 150, 220)),
            TokenKind::Context => Style::default().fg(Color::Rgb(220, 180, 120)),
            TokenKind::Note => Style::default().fg(Color::Rgb(120, 200, 160)),
            TokenKind::Interval => Style::default().fg(Color::Rgb(160, 200, 120)),
            TokenKind::Variable => Style::default().fg(Color::Rgb(200, 140, 180)),
            TokenKind::Default => Style::default().fg(Color::Rgb(200, 200, 200)),
        }
    }
}

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

fn lookup_word_kind(word: &str) -> Option<TokenKind> {
    for w in WORDS {
        if w.name == word || w.aliases.contains(&word) {
            return Some(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,
                    _ => TokenKind::Keyword,
                },
            });
        }
    }
    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 == '(' {
            let end = line.len();
            let comment_end = line[start..]
                .find(')')
                .map(|i| start + i + 1)
                .unwrap_or(end);
            tokens.push(Token {
                start,
                end: comment_end,
                kind: TokenKind::Comment,
            });
            while let Some((i, _)) = chars.peek() {
                if *i >= comment_end {
                    break;
                }
                chars.next();
            }
            continue;
        }

        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,
            });
            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 = classify_word(word);
        tokens.push(Token { start, end, kind });
    }

    tokens
}

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

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

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

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

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

    TokenKind::Default
}

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 executed_bg = Color::Rgb(40, 35, 50);
    let selected_bg = Color::Rgb(80, 60, 20);

    for token in tokens {
        if token.start > last_end {
            result.push((
                TokenKind::Default.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 is_selected {
            style = style.bg(selected_bg).add_modifier(Modifier::BOLD);
        } else if is_executed {
            style = style.bg(executed_bg);
        }

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

    if last_end < line.len() {
        result.push((TokenKind::Default.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
}