use std::collections::HashSet;
use std::sync::LazyLock;

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

use crate::model::{lookup_word, SourceSpan, WordCompile};
use crate::theme;

static EMPTY_SET: LazyLock<HashSet<String>> = LazyLock::new(HashSet::new);

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

impl TokenKind {
    pub fn style(self) -> Style {
        let theme = theme::get();
        let (fg, bg) = match self {
            TokenKind::Emit => theme.syntax.emit,
            TokenKind::Number => theme.syntax.number,
            TokenKind::String => theme.syntax.string,
            TokenKind::Comment => theme.syntax.comment,
            TokenKind::Keyword => theme.syntax.keyword,
            TokenKind::StackOp => theme.syntax.stack_op,
            TokenKind::Operator => theme.syntax.operator,
            TokenKind::Sound => theme.syntax.sound,
            TokenKind::Param => theme.syntax.param,
            TokenKind::Context => theme.syntax.context,
            TokenKind::Note => theme.syntax.note,
            TokenKind::Interval => theme.syntax.interval,
            TokenKind::Variable => theme.syntax.variable,
            TokenKind::Vary => theme.syntax.vary,
            TokenKind::Generator => theme.syntax.generator,
            TokenKind::UserDefined => theme.syntax.user_defined,
            TokenKind::Default => theme.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 {
        let theme = theme::get();
        Style::default().bg(theme.syntax.gap_bg)
    }
}

struct Token {
    start: usize,
    end: usize,
    kind: TokenKind,
    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));
    }

    let w = lookup_word(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,
        },
    };
    Some((kind, w.varargs))
}

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",
];

fn tokenize_line(line: &str, user_words: &HashSet<String>) -> 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, user_words);
        tokens.push(Token { start, end, kind, varargs });
    }

    tokens
}

fn classify_word(word: &str, user_words: &HashSet<String>) -> (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('!') || word.starts_with(',')) {
        return (TokenKind::Variable, false);
    }

    if user_words.contains(word) {
        return (TokenKind::UserDefined, false);
    }

    (TokenKind::Default, false)
}

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

pub fn highlight_line_with_runtime(
    line: &str,
    executed_spans: &[SourceSpan],
    selected_spans: &[SourceSpan],
    resolved: &[(SourceSpan, String)],
    user_words: &HashSet<String>,
) -> Vec<(Style, String, bool)> {
    let tokens = tokenize_line(line, user_words);
    let mut result = Vec::new();
    let mut last_end = 0;
    let gap_style = TokenKind::gap_style();
    let theme = theme::get();
    let annotation_style = Style::default().fg(theme.ui.text_dim);

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

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

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

        result.push((style, line[token.start..token.end].to_string(), false));

        for (span, display) in resolved {
            if token.start == span.start as usize {
                result.push((annotation_style, format!(" [{display}]"), true));
            }
        }

        last_end = token.end;
    }

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

    result
}

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