//! Tiny state-machine syntax highlighter.
//!
//! Per-line tokenisation; cross-line state is restricted to "inside a block
//! comment" / "inside a triple-quoted string".  No regexes, no tree-sitter.

use std::path::Path;

/// Languages the highlighter knows about.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Language {
    /// No highlighting (default).
    PlainText,
    /// Rust source.
    Rust,
    /// Python source.
    Python,
    /// JavaScript source.
    JavaScript,
    /// TypeScript source.
    TypeScript,
    /// C source / headers.
    C,
    /// C++ source / headers.
    Cpp,
    /// POSIX shell scripts.
    Shell,
    /// JSON documents.
    Json,
    /// TOML config files.
    Toml,
    /// Markdown / CommonMark.
    Markdown,
}

impl Language {
    /// Detect language from the file's extension (case-insensitive).
    pub fn from_path(path: &Path) -> Self {
        let ext = path
            .extension()
            .and_then(|e| e.to_str())
            .unwrap_or("")
            .to_ascii_lowercase();
        match ext.as_str() {
            "rs" => Self::Rust,
            "py" | "pyw" => Self::Python,
            "js" | "jsx" | "mjs" | "cjs" => Self::JavaScript,
            "ts" | "tsx" => Self::TypeScript,
            "c" | "h" => Self::C,
            "cc" | "cpp" | "cxx" | "hpp" | "hh" | "hxx" => Self::Cpp,
            "sh" | "bash" | "zsh" => Self::Shell,
            "json" => Self::Json,
            "toml" => Self::Toml,
            "md" | "markdown" => Self::Markdown,
            _ => Self::PlainText,
        }
    }

    /// Short label for the status bar.
    pub fn name(self) -> &'static str {
        match self {
            Self::PlainText => "text",
            Self::Rust => "rust",
            Self::Python => "python",
            Self::JavaScript => "javascript",
            Self::TypeScript => "typescript",
            Self::C => "c",
            Self::Cpp => "c++",
            Self::Shell => "shell",
            Self::Json => "json",
            Self::Toml => "toml",
            Self::Markdown => "markdown",
        }
    }
}

/// One syntactic token: a half-open byte range inside the line and its kind.
#[derive(Debug, Clone, Copy)]
pub struct Token {
    /// Byte offset of the first byte of the token.
    pub start: usize,
    /// Byte offset just past the last byte.
    pub end: usize,
    /// Token classification.
    pub kind: TokenKind,
}

/// Token classification, used by the renderer to pick colours.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenKind {
    /// Default foreground.
    Plain,
    /// Reserved word.
    Keyword,
    /// Built-in or user type.
    Type,
    /// Function call / definition name.
    Function,
    /// String / character literal.
    String,
    /// Numeric literal.
    Number,
    /// Comment.
    Comment,
    /// Punctuation / operator.
    Operator,
    /// Markdown heading run.
    Heading,
}

/// State carried over from the previous line.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LineState {
    /// Nothing special.
    Normal,
    /// Inside a `/* ... */` block comment.
    InBlockComment,
    /// Inside a Python triple-quoted string (double quotes).
    InTripleString,
}

/// Highlight one line.  Returns the tokens (in source order) and the
/// `LineState` to feed into the next line.
pub fn highlight_line(language: Language, line: &str, state: LineState) -> (Vec<Token>, LineState) {
    match language {
        Language::PlainText => (Vec::new(), LineState::Normal),
        Language::Rust => clike(line, state, RUST_KW, RUST_TYPES, true, true),
        Language::C => clike(line, state, C_KW, C_TYPES, true, true),
        Language::Cpp => clike(line, state, CPP_KW, C_TYPES, true, true),
        Language::JavaScript => clike(line, state, JS_KW, JS_TYPES, true, true),
        Language::TypeScript => clike(line, state, TS_KW, TS_TYPES, true, true),
        Language::Python => python(line, state),
        Language::Shell => shell(line),
        Language::Json => json(line),
        Language::Toml => toml(line),
        Language::Markdown => markdown(line),
    }
}

// ---------------------------------------------------------------------------
// keyword / type tables
// ---------------------------------------------------------------------------

const RUST_KW: &[&str] = &[
    "as", "async", "await", "break", "const", "continue", "crate", "dyn", "else", "enum",
    "extern", "false", "fn", "for", "if", "impl", "in", "let", "loop", "match", "mod", "move",
    "mut", "pub", "ref", "return", "self", "Self", "static", "struct", "super", "trait", "true",
    "type", "unsafe", "use", "where", "while", "yield",
];
const RUST_TYPES: &[&str] = &[
    "bool", "char", "i8", "i16", "i32", "i64", "i128", "isize", "u8", "u16", "u32", "u64", "u128",
    "usize", "f32", "f64", "str", "String", "Vec", "Option", "Result", "Box", "Rc", "Arc",
    "HashMap", "HashSet", "BTreeMap", "BTreeSet",
];

const C_KW: &[&str] = &[
    "auto", "break", "case", "const", "continue", "default", "do", "else", "enum", "extern",
    "for", "goto", "if", "inline", "register", "restrict", "return", "sizeof", "static",
    "struct", "switch", "typedef", "union", "volatile", "while", "_Alignas", "_Alignof",
    "_Atomic", "_Bool", "_Complex", "_Generic", "_Imaginary", "_Noreturn", "_Static_assert",
    "_Thread_local",
];
const C_TYPES: &[&str] = &[
    "int", "char", "void", "float", "double", "short", "long", "signed", "unsigned",
    "size_t", "ssize_t", "ptrdiff_t", "intptr_t", "uintptr_t", "uint8_t", "uint16_t",
    "uint32_t", "uint64_t", "int8_t", "int16_t", "int32_t", "int64_t", "FILE", "NULL",
    "bool", "true", "false",
];
const CPP_KW: &[&str] = &[
    "alignas", "alignof", "and", "and_eq", "asm", "auto", "bitand", "bitor", "bool", "break",
    "case", "catch", "char", "char8_t", "char16_t", "char32_t", "class", "compl", "concept",
    "const", "consteval", "constexpr", "constinit", "const_cast", "continue", "co_await",
    "co_return", "co_yield", "decltype", "default", "delete", "do", "double", "dynamic_cast",
    "else", "enum", "explicit", "export", "extern", "false", "float", "for", "friend", "goto",
    "if", "inline", "int", "long", "mutable", "namespace", "new", "noexcept", "not", "not_eq",
    "nullptr", "operator", "or", "or_eq", "private", "protected", "public", "register",
    "reinterpret_cast", "requires", "return", "short", "signed", "sizeof", "static",
    "static_assert", "static_cast", "struct", "switch", "template", "this", "thread_local",
    "throw", "true", "try", "typedef", "typeid", "typename", "union", "unsigned", "using",
    "virtual", "void", "volatile", "wchar_t", "while", "xor", "xor_eq",
];

const JS_KW: &[&str] = &[
    "async", "await", "break", "case", "catch", "class", "const", "continue", "debugger",
    "default", "delete", "do", "else", "export", "extends", "finally", "for", "from",
    "function", "if", "import", "in", "instanceof", "let", "new", "null", "of", "return",
    "static", "super", "switch", "this", "throw", "true", "false", "try", "typeof", "undefined",
    "var", "void", "while", "with", "yield",
];
const JS_TYPES: &[&str] = &["Array", "Object", "String", "Number", "Boolean", "Date", "Promise", "Map", "Set"];

const TS_KW: &[&str] = &[
    "abstract", "any", "as", "async", "await", "boolean", "break", "case", "catch", "class",
    "const", "constructor", "continue", "debugger", "declare", "default", "delete", "do", "else",
    "enum", "export", "extends", "false", "finally", "for", "from", "function", "get", "if",
    "implements", "import", "in", "infer", "instanceof", "interface", "is", "keyof", "let",
    "module", "namespace", "never", "new", "null", "number", "object", "of", "package", "private",
    "protected", "public", "readonly", "return", "set", "static", "string", "super", "switch",
    "this", "throw", "true", "try", "type", "typeof", "undefined", "unique", "unknown", "var",
    "void", "while", "with", "yield",
];
const TS_TYPES: &[&str] = &[
    "Array", "Object", "String", "Number", "Boolean", "Date", "Promise", "Map", "Set",
    "Record", "Partial", "Readonly", "Pick", "Omit",
];

const PY_KW: &[&str] = &[
    "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class",
    "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return",
    "try", "while", "with", "yield", "match", "case",
];
const PY_TYPES: &[&str] = &[
    "int", "str", "float", "bool", "bytes", "list", "tuple", "dict", "set", "frozenset",
    "object", "type", "complex", "range",
];

const SHELL_KW: &[&str] = &[
    "if", "then", "else", "elif", "fi", "case", "esac", "for", "while", "until", "do", "done",
    "in", "function", "select", "return", "exit", "break", "continue", "local", "readonly",
    "export", "unset", "shift", "set", "test", "trap", "alias", "unalias", "source",
];

const JSON_KW: &[&str] = &["true", "false", "null"];

// ---------------------------------------------------------------------------
// C-like tokenizer (Rust / C / C++ / JS / TS)
// ---------------------------------------------------------------------------

fn clike(
    line: &str,
    state: LineState,
    keywords: &[&str],
    types: &[&str],
    line_comment_dslash: bool,
    block_comment_slash_star: bool,
) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let mut i = 0usize;
    let mut cur = state;

    if cur == LineState::InBlockComment {
        match line.find("*/") {
            Some(end) => {
                tokens.push(Token { start: 0, end: end + 2, kind: TokenKind::Comment });
                i = end + 2;
                cur = LineState::Normal;
            }
            None => {
                tokens.push(Token { start: 0, end: bytes.len(), kind: TokenKind::Comment });
                return (tokens, LineState::InBlockComment);
            }
        }
    }

    while i < bytes.len() {
        let b = bytes[i];
        if b.is_ascii_whitespace() {
            i += 1;
            continue;
        }
        if block_comment_slash_star && b == b'/' && bytes.get(i + 1) == Some(&b'*') {
            let start = i;
            match line[i + 2..].find("*/") {
                Some(end) => {
                    let abs = i + 2 + end + 2;
                    tokens.push(Token { start, end: abs, kind: TokenKind::Comment });
                    i = abs;
                }
                None => {
                    tokens.push(Token { start, end: bytes.len(), kind: TokenKind::Comment });
                    return (tokens, LineState::InBlockComment);
                }
            }
            continue;
        }
        if line_comment_dslash && b == b'/' && bytes.get(i + 1) == Some(&b'/') {
            tokens.push(Token { start: i, end: bytes.len(), kind: TokenKind::Comment });
            return (tokens, cur);
        }
        if b == b'"' {
            i = scan_string(&mut tokens, line, i, b'"');
            continue;
        }
        if b == b'\'' {
            i = scan_string(&mut tokens, line, i, b'\'');
            continue;
        }
        if b.is_ascii_digit() {
            let start = i;
            while i < bytes.len()
                && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'.' || bytes[i] == b'_')
            {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Number });
            continue;
        }
        if is_ident_start(b) {
            let start = i;
            while i < bytes.len() && is_ident_cont(bytes[i]) {
                i += 1;
            }
            let word = &line[start..i];
            let kind = classify_word(word, keywords, types, bytes, i);
            tokens.push(Token { start, end: i, kind });
            continue;
        }
        let start = i;
        i += 1;
        tokens.push(Token { start, end: i, kind: TokenKind::Operator });
    }
    (tokens, cur)
}

fn classify_word(word: &str, keywords: &[&str], types: &[&str], bytes: &[u8], end: usize) -> TokenKind {
    if keywords.contains(&word) {
        return TokenKind::Keyword;
    }
    if types.contains(&word) {
        return TokenKind::Type;
    }
    if word.starts_with(|c: char| c.is_ascii_uppercase()) {
        return TokenKind::Type;
    }
    let mut j = end;
    while j < bytes.len() && bytes[j] == b' ' {
        j += 1;
    }
    if bytes.get(j) == Some(&b'(') {
        return TokenKind::Function;
    }
    TokenKind::Plain
}

fn scan_string(tokens: &mut Vec<Token>, line: &str, start: usize, q: u8) -> usize {
    let bytes = line.as_bytes();
    let mut i = start + 1;
    while i < bytes.len() {
        if bytes[i] == b'\\' && i + 1 < bytes.len() {
            i += 2;
            continue;
        }
        if bytes[i] == q {
            i += 1;
            break;
        }
        i += 1;
    }
    tokens.push(Token { start, end: i, kind: TokenKind::String });
    i
}

fn is_ident_start(b: u8) -> bool {
    b == b'_' || b.is_ascii_alphabetic()
}

fn is_ident_cont(b: u8) -> bool {
    b == b'_' || b.is_ascii_alphanumeric()
}

// ---------------------------------------------------------------------------
// Python
// ---------------------------------------------------------------------------

fn python(line: &str, state: LineState) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let mut i = 0usize;
    let mut cur = state;

    if cur == LineState::InTripleString {
        match line.find("\"\"\"") {
            Some(end) => {
                tokens.push(Token { start: 0, end: end + 3, kind: TokenKind::String });
                i = end + 3;
                cur = LineState::Normal;
            }
            None => {
                tokens.push(Token { start: 0, end: bytes.len(), kind: TokenKind::String });
                return (tokens, LineState::InTripleString);
            }
        }
    }

    while i < bytes.len() {
        let b = bytes[i];
        if b.is_ascii_whitespace() {
            i += 1;
            continue;
        }
        if b == b'#' {
            tokens.push(Token { start: i, end: bytes.len(), kind: TokenKind::Comment });
            return (tokens, cur);
        }
        if i + 2 < bytes.len() && &line[i..i + 3] == "\"\"\"" {
            let start = i;
            match line[i + 3..].find("\"\"\"") {
                Some(end) => {
                    let abs = i + 3 + end + 3;
                    tokens.push(Token { start, end: abs, kind: TokenKind::String });
                    i = abs;
                }
                None => {
                    tokens.push(Token { start, end: bytes.len(), kind: TokenKind::String });
                    return (tokens, LineState::InTripleString);
                }
            }
            continue;
        }
        if b == b'"' || b == b'\'' {
            i = scan_string(&mut tokens, line, i, b);
            continue;
        }
        if b.is_ascii_digit() {
            let start = i;
            while i < bytes.len()
                && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'.' || bytes[i] == b'_')
            {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Number });
            continue;
        }
        if is_ident_start(b) {
            let start = i;
            while i < bytes.len() && is_ident_cont(bytes[i]) {
                i += 1;
            }
            let word = &line[start..i];
            let kind = classify_word(word, PY_KW, PY_TYPES, bytes, i);
            tokens.push(Token { start, end: i, kind });
            continue;
        }
        let start = i;
        i += 1;
        tokens.push(Token { start, end: i, kind: TokenKind::Operator });
    }
    (tokens, cur)
}

// ---------------------------------------------------------------------------
// Shell
// ---------------------------------------------------------------------------

fn shell(line: &str) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let mut i = 0;
    while i < bytes.len() {
        let b = bytes[i];
        if b.is_ascii_whitespace() {
            i += 1;
            continue;
        }
        if b == b'#' {
            tokens.push(Token { start: i, end: bytes.len(), kind: TokenKind::Comment });
            return (tokens, LineState::Normal);
        }
        if b == b'"' || b == b'\'' {
            i = scan_string(&mut tokens, line, i, b);
            continue;
        }
        if b == b'$' {
            let start = i;
            i += 1;
            while i < bytes.len() && (is_ident_cont(bytes[i]) || bytes[i] == b'{' || bytes[i] == b'}') {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Type });
            continue;
        }
        if b.is_ascii_digit() {
            let start = i;
            while i < bytes.len() && bytes[i].is_ascii_alphanumeric() {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Number });
            continue;
        }
        if is_ident_start(b) {
            let start = i;
            while i < bytes.len() && (is_ident_cont(bytes[i]) || bytes[i] == b'-') {
                i += 1;
            }
            let word = &line[start..i];
            let kind = if SHELL_KW.contains(&word) {
                TokenKind::Keyword
            } else {
                let mut j = i;
                while j < bytes.len() && bytes[j] == b' ' {
                    j += 1;
                }
                if bytes.get(j) == Some(&b'(') {
                    TokenKind::Function
                } else {
                    TokenKind::Plain
                }
            };
            tokens.push(Token { start, end: i, kind });
            continue;
        }
        let start = i;
        i += 1;
        tokens.push(Token { start, end: i, kind: TokenKind::Operator });
    }
    (tokens, LineState::Normal)
}

// ---------------------------------------------------------------------------
// JSON
// ---------------------------------------------------------------------------

fn json(line: &str) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let mut i = 0;
    while i < bytes.len() {
        let b = bytes[i];
        if b.is_ascii_whitespace() {
            i += 1;
            continue;
        }
        if b == b'"' {
            i = scan_string(&mut tokens, line, i, b'"');
            continue;
        }
        if b.is_ascii_digit() || b == b'-' {
            let start = i;
            i += 1;
            while i < bytes.len()
                && (bytes[i].is_ascii_digit() || bytes[i] == b'.' || bytes[i] == b'e' || bytes[i] == b'E' || bytes[i] == b'+' || bytes[i] == b'-')
            {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Number });
            continue;
        }
        if is_ident_start(b) {
            let start = i;
            while i < bytes.len() && is_ident_cont(bytes[i]) {
                i += 1;
            }
            let word = &line[start..i];
            let kind = if JSON_KW.contains(&word) {
                TokenKind::Keyword
            } else {
                TokenKind::Plain
            };
            tokens.push(Token { start, end: i, kind });
            continue;
        }
        let start = i;
        i += 1;
        tokens.push(Token { start, end: i, kind: TokenKind::Operator });
    }
    (tokens, LineState::Normal)
}

// ---------------------------------------------------------------------------
// TOML
// ---------------------------------------------------------------------------

fn toml(line: &str) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let trimmed_start = line.bytes().position(|c| !c.is_ascii_whitespace()).unwrap_or(bytes.len());
    if trimmed_start < bytes.len() && bytes[trimmed_start] == b'#' {
        tokens.push(Token { start: trimmed_start, end: bytes.len(), kind: TokenKind::Comment });
        return (tokens, LineState::Normal);
    }
    if trimmed_start < bytes.len() && bytes[trimmed_start] == b'[' {
        let end = line[trimmed_start..]
            .find(']')
            .map(|p| trimmed_start + p + 1)
            .unwrap_or(bytes.len());
        tokens.push(Token { start: trimmed_start, end, kind: TokenKind::Heading });
        return (tokens, LineState::Normal);
    }
    let mut i = 0;
    let mut seen_eq = false;
    while i < bytes.len() {
        let b = bytes[i];
        if b.is_ascii_whitespace() {
            i += 1;
            continue;
        }
        if b == b'#' {
            tokens.push(Token { start: i, end: bytes.len(), kind: TokenKind::Comment });
            return (tokens, LineState::Normal);
        }
        if b == b'"' || b == b'\'' {
            i = scan_string(&mut tokens, line, i, b);
            continue;
        }
        if b == b'=' {
            tokens.push(Token { start: i, end: i + 1, kind: TokenKind::Operator });
            seen_eq = true;
            i += 1;
            continue;
        }
        if b.is_ascii_digit() || (b == b'-' && seen_eq) {
            let start = i;
            i += 1;
            while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'.' || bytes[i] == b'_' || bytes[i] == b'-') {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::Number });
            continue;
        }
        if is_ident_start(b) {
            let start = i;
            while i < bytes.len() && (is_ident_cont(bytes[i]) || bytes[i] == b'-' || bytes[i] == b'.') {
                i += 1;
            }
            let word = &line[start..i];
            let kind = if !seen_eq {
                TokenKind::Type
            } else if word == "true" || word == "false" {
                TokenKind::Keyword
            } else {
                TokenKind::Plain
            };
            tokens.push(Token { start, end: i, kind });
            continue;
        }
        let start = i;
        i += 1;
        tokens.push(Token { start, end: i, kind: TokenKind::Operator });
    }
    (tokens, LineState::Normal)
}

// ---------------------------------------------------------------------------
// Markdown
// ---------------------------------------------------------------------------

fn markdown(line: &str) -> (Vec<Token>, LineState) {
    let bytes = line.as_bytes();
    let mut tokens = Vec::new();
    let leading = line.bytes().take_while(|b| *b == b' ').count();
    let rest = &line[leading..];
    if rest.starts_with('#') {
        tokens.push(Token { start: 0, end: bytes.len(), kind: TokenKind::Heading });
        return (tokens, LineState::Normal);
    }
    if rest.starts_with("```") {
        tokens.push(Token { start: 0, end: bytes.len(), kind: TokenKind::Keyword });
        return (tokens, LineState::Normal);
    }
    if rest.starts_with("> ") {
        tokens.push(Token { start: leading, end: bytes.len(), kind: TokenKind::Comment });
        return (tokens, LineState::Normal);
    }
    let mut i = 0usize;
    while i < bytes.len() {
        let b = bytes[i];
        if b == b'`' {
            let start = i;
            i += 1;
            while i < bytes.len() && bytes[i] != b'`' {
                i += 1;
            }
            if i < bytes.len() {
                i += 1;
            }
            tokens.push(Token { start, end: i, kind: TokenKind::String });
            continue;
        }
        if b == b'*' || b == b'_' {
            let start = i;
            let mark = b;
            let count = if bytes.get(i + 1) == Some(&mark) { 2 } else { 1 };
            i += count;
            while i < bytes.len() && bytes[i] != mark {
                i += 1;
            }
            i = (i + count).min(bytes.len());
            tokens.push(Token { start, end: i, kind: TokenKind::Keyword });
            continue;
        }
        i += 1;
    }
    (tokens, LineState::Normal)
}
