// lexer_lua.c -- Lua language tokenizer
//
// Same pattern as lexer_go.c: a get_next_token() loop that advances a cursor
// through the source, producing Token structs.  After each token we paint
// the per-byte token-type array.  A second pass marks identifiers before '('
// as functions.

#include "lexer/lexer.h"
#include <ctype.h>
#include <string.h>

////////////////////////////////
// Lua keyword tables

typedef struct LuaKeywordEntry {
    const char *word;
    Token_Type  type;
} LuaKeywordEntry;

static const LuaKeywordEntry lua_keywords[] = {
    // Keywords
    {"and",         TOK_KEYWORD},
    {"break",       TOK_KEYWORD},
    {"do",          TOK_KEYWORD},
    {"else",        TOK_KEYWORD},
    {"elseif",      TOK_KEYWORD},
    {"end",         TOK_KEYWORD},
    {"for",         TOK_KEYWORD},
    {"function",    TOK_KEYWORD},
    {"goto",        TOK_KEYWORD},
    {"if",          TOK_KEYWORD},
    {"in",          TOK_KEYWORD},
    {"local",       TOK_KEYWORD},
    {"not",         TOK_KEYWORD},
    {"or",          TOK_KEYWORD},
    {"repeat",      TOK_KEYWORD},
    {"return",      TOK_KEYWORD},
    {"then",        TOK_KEYWORD},
    {"until",       TOK_KEYWORD},
    {"while",       TOK_KEYWORD},

    // Values
    {"true",        TOK_VALUE},
    {"false",       TOK_VALUE},
    {"nil",         TOK_VALUE},

    // Built-in globals / types (use TOK_TYPE for distinct coloring)
    {"self",        TOK_TYPE},

    // Built-in functions (TOK_MODIFIER for distinct coloring)
    {"assert",      TOK_MODIFIER},
    {"collectgarbage", TOK_MODIFIER},
    {"dofile",      TOK_MODIFIER},
    {"error",       TOK_MODIFIER},
    {"getmetatable", TOK_MODIFIER},
    {"ipairs",      TOK_MODIFIER},
    {"load",        TOK_MODIFIER},
    {"loadfile",    TOK_MODIFIER},
    {"next",        TOK_MODIFIER},
    {"pairs",       TOK_MODIFIER},
    {"pcall",       TOK_MODIFIER},
    {"print",       TOK_MODIFIER},
    {"rawequal",    TOK_MODIFIER},
    {"rawget",      TOK_MODIFIER},
    {"rawlen",      TOK_MODIFIER},
    {"rawset",      TOK_MODIFIER},
    {"require",     TOK_MODIFIER},
    {"select",      TOK_MODIFIER},
    {"setmetatable", TOK_MODIFIER},
    {"tonumber",    TOK_MODIFIER},
    {"tostring",    TOK_MODIFIER},
    {"type",        TOK_MODIFIER},
    {"unpack",      TOK_MODIFIER},
    {"xpcall",      TOK_MODIFIER},
};

#define LUA_KEYWORD_COUNT (S32)(sizeof(lua_keywords) / sizeof(lua_keywords[0]))

static Token_Type lua_lookup_keyword(const char *word, S32 len) {
    for (S32 i = 0; i < LUA_KEYWORD_COUNT; i++) {
        const char *kw = lua_keywords[i].word;
        S32 kwlen = (S32)strlen(kw);
        if (kwlen == len && memcmp(kw, word, len) == 0)
            return lua_keywords[i].type;
    }
    return TOK_IDENTIFIER;
}

////////////////////////////////
// Character helpers

static B32 lua_is_ident_start(char c) {
    return isalpha((unsigned char)c) || c == '_';
}

static B32 lua_is_ident_char(char c) {
    return isalnum((unsigned char)c) || c == '_';
}

static B32 lua_is_hex(char c) {
    return isdigit((unsigned char)c) ||
           (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F');
}

////////////////////////////////
// Long bracket level detection
// Returns the level (number of '=' signs) if we're at a long bracket opening
// like [[ or [==[ etc. Returns -1 if not a long bracket.

static S32 lua_long_bracket_level(const char *p, const char *max_p) {
    if (p >= max_p || *p != '[') return -1;
    p++;
    S32 level = 0;
    while (p < max_p && *p == '=') { level++; p++; }
    if (p < max_p && *p == '[') return level;
    return -1;
}

// Scan past the closing long bracket of the given level.
// Cursor should be right after the opening [=*[.
static void lua_scan_long_bracket_close(Tokenizer *tok, S32 level) {
    while (tok->t < tok->max_t) {
        if (*tok->t == ']') {
            const char *p = tok->t + 1;
            S32 count = 0;
            while (p < tok->max_t && *p == '=' && count < level) { count++; p++; }
            if (count == level && p < tok->max_t && *p == ']') {
                tok->t = p + 1;
                return;
            }
        }
        tok->t++;
    }
}

////////////////////////////////
// Individual token parsers

static Token lua_parse_identifier(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    const char *begin = tok->t;
    while (tok->t < tok->max_t && lua_is_ident_char(*tok->t))
        tok->t++;
    S32 len = (S32)(tok->t - begin);
    token.type = lua_lookup_keyword(begin, len);
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_number(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_NUMBER;

    char c = *tok->t;
    tok->t++;

    if (c == '0' && tok->t < tok->max_t && (*tok->t == 'x' || *tok->t == 'X')) {
        // Hex literal
        tok->t++;
        while (tok->t < tok->max_t && lua_is_hex(*tok->t)) tok->t++;
        // Hex float
        if (tok->t < tok->max_t && *tok->t == '.') {
            tok->t++;
            while (tok->t < tok->max_t && lua_is_hex(*tok->t)) tok->t++;
        }
        // Hex exponent (p/P)
        if (tok->t < tok->max_t && (*tok->t == 'p' || *tok->t == 'P')) {
            tok->t++;
            if (tok->t < tok->max_t && (*tok->t == '+' || *tok->t == '-')) tok->t++;
            while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
        }
    } else {
        // Decimal
        while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
        // Float
        if (tok->t < tok->max_t && *tok->t == '.') {
            if ((tok->t + 1) < tok->max_t && isdigit((unsigned char)*(tok->t + 1))) {
                tok->t++;
                while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
            } else if ((tok->t + 1) >= tok->max_t || !lua_is_ident_start(*(tok->t + 1))) {
                tok->t++;
            }
        }
        // Exponent
        if (tok->t < tok->max_t && (*tok->t == 'e' || *tok->t == 'E')) {
            tok->t++;
            if (tok->t < tok->max_t && (*tok->t == '+' || *tok->t == '-')) tok->t++;
            while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
        }
    }

    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_dot_number(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_NUMBER;

    tok->t++; // skip '.'
    while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
    if (tok->t < tok->max_t && (*tok->t == 'e' || *tok->t == 'E')) {
        tok->t++;
        if (tok->t < tok->max_t && (*tok->t == '+' || *tok->t == '-')) tok->t++;
        while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
    }

    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_string(Tokenizer *tok, char quote) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_STRING_LITERAL;

    B32 escape = 0;
    tok->t++; // skip opening quote
    while (tok->t < tok->max_t && *tok->t != '\n') {
        if (*tok->t == quote && !escape) { tok->t++; break; }
        escape = !escape && (*tok->t == '\\');
        tok->t++;
    }
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_long_string(Tokenizer *tok, S32 level) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_STRING_LITERAL;

    // Skip past opening [=*[
    tok->t++; // first [
    tok->t += level; // = signs
    tok->t++; // second [

    lua_scan_long_bracket_close(tok, level);

    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_comment(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);

    tok->t += 2; // skip '--'

    // Check for long comment --[=*[
    S32 level = lua_long_bracket_level(tok->t, tok->max_t);
    if (level >= 0) {
        token.type = TOK_MULTILINE_COMMENT;
        tok->t++; // first [
        tok->t += level; // = signs
        tok->t++; // second [
        lua_scan_long_bracket_close(tok, level);
    } else {
        // Single-line comment
        token.type = TOK_COMMENT;
        while (tok->t < tok->max_t && *tok->t != '\n') tok->t++;
    }

    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

static Token lua_parse_operator(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_OPERATION;

    char c = *tok->t;
    tok->t++;

    if (tok->t < tok->max_t) {
        char n = *tok->t;
        // Two-character operators
        if ((c == '=' && n == '=') || (c == '~' && n == '=') ||
            (c == '<' && n == '=') || (c == '>' && n == '=') ||
            (c == '.' && n == '.') || (c == '<' && n == '<') ||
            (c == '>' && n == '>') || (c == '/' && n == '/')) {
            tok->t++;
            // Three-character: ..  can be followed by . to make ...
            if (c == '.' && n == '.' && tok->t < tok->max_t && *tok->t == '.') {
                tok->t++;
            }
        }
    }

    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

////////////////////////////////
// Main get_next_token

static Token lua_get_next_token(Tokenizer *tok) {
    tokenizer_eat_whitespace(tok);

    Token token = {0};
    token.start = (S32)(tok->t - tok->buf);
    token.type = TOK_DEFAULT;

    if (tok->t >= tok->max_t) {
        token.len = 0;
        return token;
    }

    tok->start_t = tok->t;
    char c = *tok->t;

    if (lua_is_ident_start(c)) {
        return lua_parse_identifier(tok);
    }
    if (isdigit((unsigned char)c)) {
        return lua_parse_number(tok);
    }

    switch (c) {
    case '\'':
    case '"':
        return lua_parse_string(tok, c);

    case '[': {
        // Check for long string [=*[
        S32 level = lua_long_bracket_level(tok->t, tok->max_t);
        if (level >= 0) {
            return lua_parse_long_string(tok, level);
        }
        // Plain bracket punctuation
        token.type = TOK_PUNCTUATION;
        tok->t++;
        token.len = 1;
        return token;
    }

    case '-':
        // Check for comment --
        if ((tok->t + 1) < tok->max_t && *(tok->t + 1) == '-') {
            return lua_parse_comment(tok);
        }
        // Minus operator
        token.type = TOK_OPERATION;
        tok->t++;
        token.len = 1;
        return token;

    case '.':
        // .digits = float literal
        if ((tok->t + 1) < tok->max_t && isdigit((unsigned char)*(tok->t + 1))) {
            return lua_parse_dot_number(tok);
        }
        // .. or ... or just .
        return lua_parse_operator(tok);

    // Punctuation
    case ';': case ',': case ':':
    case '{': case '}': case '(': case ')': case ']':
        token.type = TOK_PUNCTUATION;
        tok->t++;
        token.len = 1;
        return token;

    // Directive: shebang or labels (::label::)
    case '#':
        // Shebang line or length operator
        if (tok->t == tok->buf && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '!') {
            token.type = TOK_DIRECTIVE;
            while (tok->t < tok->max_t && *tok->t != '\n') tok->t++;
            token.len = (S32)(tok->t - tok->start_t);
            return token;
        }
        token.type = TOK_OPERATION;
        tok->t++;
        token.len = 1;
        return token;

    // Operators
    case '=': case '~': case '<': case '>':
    case '+': case '*': case '/': case '%':
    case '^': case '&': case '|':
        return lua_parse_operator(tok);

    default:
        token.type = TOK_INVALID;
        tok->t++;
        token.len = 1;
        return token;
    }
}

////////////////////////////////
// tokenize_lua -- main entry point

static void tokenize_lua(const char *data, S32 length, U8 *out_tokens) {
    memset(out_tokens, TOK_DEFAULT, length);

    Tokenizer tok;
    tokenizer_init(&tok, data, length);

    Token prev = {0};
    prev.type = TOK_DEFAULT;

    while (tok.t < tok.max_t) {
        Token token = lua_get_next_token(&tok);
        if (token.len == 0) break;

        paint_token(out_tokens, token.start, token.len, token.type);

        // Retroactively mark identifier before '(' as a function
        if (token.type == TOK_PUNCTUATION && token.len == 1 &&
            data[token.start] == '(' && prev.type == TOK_IDENTIFIER) {
            paint_token(out_tokens, prev.start, prev.len, TOK_FUNCTION);
        }

        prev = token;
    }
}