// lexer_sql.c -- SQL 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.
//
// SQL keywords are case-insensitive, so comparisons use a case-insensitive
// match.  Covers standard SQL plus common extensions (MySQL, PostgreSQL, etc.).

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

////////////////////////////////
// SQL keyword tables

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

static const SQLKeywordEntry sql_keywords[] = {
    // DML / DQL keywords
    {"SELECT",      TOK_KEYWORD},
    {"FROM",        TOK_KEYWORD},
    {"WHERE",       TOK_KEYWORD},
    {"INSERT",      TOK_KEYWORD},
    {"INTO",        TOK_KEYWORD},
    {"UPDATE",      TOK_KEYWORD},
    {"DELETE",      TOK_KEYWORD},
    {"SET",         TOK_KEYWORD},
    {"VALUES",      TOK_KEYWORD},
    {"AS",          TOK_KEYWORD},
    {"ON",          TOK_KEYWORD},
    {"JOIN",        TOK_KEYWORD},
    {"INNER",       TOK_KEYWORD},
    {"LEFT",        TOK_KEYWORD},
    {"RIGHT",       TOK_KEYWORD},
    {"OUTER",       TOK_KEYWORD},
    {"FULL",        TOK_KEYWORD},
    {"CROSS",       TOK_KEYWORD},
    {"NATURAL",     TOK_KEYWORD},
    {"USING",       TOK_KEYWORD},
    {"ORDER",       TOK_KEYWORD},
    {"BY",          TOK_KEYWORD},
    {"GROUP",       TOK_KEYWORD},
    {"HAVING",      TOK_KEYWORD},
    {"LIMIT",       TOK_KEYWORD},
    {"OFFSET",      TOK_KEYWORD},
    {"UNION",       TOK_KEYWORD},
    {"ALL",         TOK_KEYWORD},
    {"DISTINCT",    TOK_KEYWORD},
    {"TOP",         TOK_KEYWORD},
    {"FETCH",       TOK_KEYWORD},
    {"NEXT",        TOK_KEYWORD},
    {"ROWS",        TOK_KEYWORD},
    {"ONLY",        TOK_KEYWORD},
    {"FIRST",       TOK_KEYWORD},
    {"LAST",        TOK_KEYWORD},

    // DDL keywords
    {"CREATE",      TOK_KEYWORD},
    {"ALTER",       TOK_KEYWORD},
    {"DROP",        TOK_KEYWORD},
    {"TABLE",       TOK_KEYWORD},
    {"VIEW",        TOK_KEYWORD},
    {"INDEX",       TOK_KEYWORD},
    {"DATABASE",    TOK_KEYWORD},
    {"SCHEMA",      TOK_KEYWORD},
    {"COLUMN",      TOK_KEYWORD},
    {"ADD",         TOK_KEYWORD},
    {"RENAME",      TOK_KEYWORD},
    {"TRUNCATE",    TOK_KEYWORD},
    {"REPLACE",     TOK_KEYWORD},
    {"TEMPORARY",   TOK_KEYWORD},
    {"TEMP",        TOK_KEYWORD},
    {"IF",          TOK_KEYWORD},
    {"EXISTS",      TOK_KEYWORD},
    {"CASCADE",     TOK_KEYWORD},
    {"RESTRICT",    TOK_KEYWORD},

    // Constraints & keys
    {"PRIMARY",     TOK_KEYWORD},
    {"KEY",         TOK_KEYWORD},
    {"FOREIGN",     TOK_KEYWORD},
    {"REFERENCES",  TOK_KEYWORD},
    {"UNIQUE",      TOK_KEYWORD},
    {"CHECK",       TOK_KEYWORD},
    {"CONSTRAINT",  TOK_KEYWORD},
    {"DEFAULT",     TOK_KEYWORD},
    {"AUTOINCREMENT", TOK_KEYWORD},
    {"AUTO_INCREMENT", TOK_KEYWORD},
    {"IDENTITY",    TOK_KEYWORD},

    // Control flow / procedural
    {"BEGIN",       TOK_KEYWORD},
    {"END",         TOK_KEYWORD},
    {"COMMIT",      TOK_KEYWORD},
    {"ROLLBACK",    TOK_KEYWORD},
    {"SAVEPOINT",   TOK_KEYWORD},
    {"TRANSACTION", TOK_KEYWORD},
    {"RETURN",      TOK_KEYWORD},
    {"RETURNS",     TOK_KEYWORD},
    {"DECLARE",     TOK_KEYWORD},
    {"CURSOR",      TOK_KEYWORD},
    {"OPEN",        TOK_KEYWORD},
    {"CLOSE",       TOK_KEYWORD},
    {"DEALLOCATE",  TOK_KEYWORD},
    {"EXEC",        TOK_KEYWORD},
    {"EXECUTE",     TOK_KEYWORD},
    {"CALL",        TOK_KEYWORD},
    {"PROCEDURE",   TOK_KEYWORD},
    {"FUNCTION",    TOK_KEYWORD},
    {"TRIGGER",     TOK_KEYWORD},
    {"CASE",        TOK_KEYWORD},
    {"WHEN",        TOK_KEYWORD},
    {"THEN",        TOK_KEYWORD},
    {"ELSE",        TOK_KEYWORD},
    {"WHILE",       TOK_KEYWORD},
    {"LOOP",        TOK_KEYWORD},
    {"FOR",         TOK_KEYWORD},
    {"EACH",        TOK_KEYWORD},
    {"ROW",         TOK_KEYWORD},
    {"AFTER",       TOK_KEYWORD},
    {"BEFORE",      TOK_KEYWORD},
    {"INSTEAD",     TOK_KEYWORD},
    {"OF",          TOK_KEYWORD},

    // Logical operators (keywords)
    {"AND",         TOK_KEYWORD},
    {"OR",          TOK_KEYWORD},
    {"NOT",         TOK_KEYWORD},
    {"IN",          TOK_KEYWORD},
    {"BETWEEN",     TOK_KEYWORD},
    {"LIKE",        TOK_KEYWORD},
    {"ILIKE",       TOK_KEYWORD},
    {"IS",          TOK_KEYWORD},
    {"ANY",         TOK_KEYWORD},
    {"SOME",        TOK_KEYWORD},
    {"EXCEPT",      TOK_KEYWORD},
    {"INTERSECT",   TOK_KEYWORD},
    {"WITH",        TOK_KEYWORD},
    {"RECURSIVE",   TOK_KEYWORD},
    {"OVER",        TOK_KEYWORD},
    {"PARTITION",   TOK_KEYWORD},
    {"WINDOW",      TOK_KEYWORD},
    {"LATERAL",     TOK_KEYWORD},

    // Misc
    {"GRANT",       TOK_KEYWORD},
    {"REVOKE",      TOK_KEYWORD},
    {"EXPLAIN",     TOK_KEYWORD},
    {"ANALYZE",     TOK_KEYWORD},
    {"VACUUM",      TOK_KEYWORD},
    {"PRAGMA",      TOK_KEYWORD},
    {"DESCRIBE",    TOK_KEYWORD},
    {"SHOW",        TOK_KEYWORD},
    {"USE",         TOK_KEYWORD},
    {"COPY",        TOK_KEYWORD},
    {"PERFORM",     TOK_KEYWORD},
    {"RAISE",       TOK_KEYWORD},

    // Modifiers
    {"ASC",         TOK_MODIFIER},
    {"DESC",        TOK_MODIFIER},
    {"NULLS",       TOK_MODIFIER},
    {"NOT",         TOK_MODIFIER},

    // Data types
    {"INT",         TOK_TYPE},
    {"INTEGER",     TOK_TYPE},
    {"SMALLINT",    TOK_TYPE},
    {"BIGINT",      TOK_TYPE},
    {"TINYINT",     TOK_TYPE},
    {"MEDIUMINT",   TOK_TYPE},
    {"SERIAL",      TOK_TYPE},
    {"BIGSERIAL",   TOK_TYPE},
    {"FLOAT",       TOK_TYPE},
    {"REAL",        TOK_TYPE},
    {"DOUBLE",      TOK_TYPE},
    {"DECIMAL",     TOK_TYPE},
    {"NUMERIC",     TOK_TYPE},
    {"PRECISION",   TOK_TYPE},
    {"CHAR",        TOK_TYPE},
    {"VARCHAR",     TOK_TYPE},
    {"TEXT",        TOK_TYPE},
    {"NCHAR",       TOK_TYPE},
    {"NVARCHAR",    TOK_TYPE},
    {"NTEXT",       TOK_TYPE},
    {"BLOB",        TOK_TYPE},
    {"CLOB",        TOK_TYPE},
    {"BYTEA",       TOK_TYPE},
    {"BOOLEAN",     TOK_TYPE},
    {"BOOL",        TOK_TYPE},
    {"DATE",        TOK_TYPE},
    {"TIME",        TOK_TYPE},
    {"DATETIME",    TOK_TYPE},
    {"TIMESTAMP",   TOK_TYPE},
    {"TIMESTAMPTZ", TOK_TYPE},
    {"INTERVAL",    TOK_TYPE},
    {"UUID",        TOK_TYPE},
    {"JSON",        TOK_TYPE},
    {"JSONB",       TOK_TYPE},
    {"XML",         TOK_TYPE},
    {"ARRAY",       TOK_TYPE},
    {"ENUM",        TOK_TYPE},
    {"MONEY",       TOK_TYPE},
    {"BIT",         TOK_TYPE},
    {"VARBIT",      TOK_TYPE},
    {"INET",        TOK_TYPE},
    {"CIDR",        TOK_TYPE},
    {"MACADDR",     TOK_TYPE},
    {"POINT",       TOK_TYPE},
    {"LINE",        TOK_TYPE},
    {"POLYGON",     TOK_TYPE},
    {"GEOMETRY",    TOK_TYPE},
    {"GEOGRAPHY",   TOK_TYPE},

    // Values
    {"NULL",        TOK_VALUE},
    {"TRUE",        TOK_VALUE},
    {"FALSE",       TOK_VALUE},
    {"CURRENT_DATE",     TOK_VALUE},
    {"CURRENT_TIME",     TOK_VALUE},
    {"CURRENT_TIMESTAMP", TOK_VALUE},
    {"CURRENT_USER",     TOK_VALUE},

    // Built-in aggregate / window / common functions (TOK_MODIFIER)
    {"COUNT",       TOK_MODIFIER},
    {"SUM",         TOK_MODIFIER},
    {"AVG",         TOK_MODIFIER},
    {"MIN",         TOK_MODIFIER},
    {"MAX",         TOK_MODIFIER},
    {"COALESCE",    TOK_MODIFIER},
    {"NULLIF",      TOK_MODIFIER},
    {"CAST",        TOK_MODIFIER},
    {"CONVERT",     TOK_MODIFIER},
    {"IFNULL",      TOK_MODIFIER},
    {"ISNULL",      TOK_MODIFIER},
    {"NVL",         TOK_MODIFIER},
    {"ROW_NUMBER",  TOK_MODIFIER},
    {"RANK",        TOK_MODIFIER},
    {"DENSE_RANK",  TOK_MODIFIER},
    {"NTILE",       TOK_MODIFIER},
    {"LAG",         TOK_MODIFIER},
    {"LEAD",        TOK_MODIFIER},
    {"FIRST_VALUE", TOK_MODIFIER},
    {"LAST_VALUE",  TOK_MODIFIER},
    {"SUBSTR",      TOK_MODIFIER},
    {"SUBSTRING",   TOK_MODIFIER},
    {"TRIM",        TOK_MODIFIER},
    {"UPPER",       TOK_MODIFIER},
    {"LOWER",       TOK_MODIFIER},
    {"LENGTH",      TOK_MODIFIER},
    {"CONCAT",      TOK_MODIFIER},
    {"REPLACE",     TOK_MODIFIER},
    {"ABS",         TOK_MODIFIER},
    {"ROUND",       TOK_MODIFIER},
    {"CEIL",        TOK_MODIFIER},
    {"FLOOR",       TOK_MODIFIER},
    {"NOW",         TOK_MODIFIER},
    {"EXTRACT",     TOK_MODIFIER},
    {"STRING_AGG",  TOK_MODIFIER},
    {"GROUP_CONCAT", TOK_MODIFIER},
    {"ARRAY_AGG",   TOK_MODIFIER},
    {"GREATEST",    TOK_MODIFIER},
    {"LEAST",       TOK_MODIFIER},
};

#define SQL_KEYWORD_COUNT (S32)(sizeof(sql_keywords) / sizeof(sql_keywords[0]))

// Case-insensitive keyword lookup
static Token_Type sql_lookup_keyword(const char *word, S32 len) {
    for (S32 i = 0; i < SQL_KEYWORD_COUNT; i++) {
        const char *kw = sql_keywords[i].word;
        S32 kwlen = (S32)strlen(kw);
        if (kwlen != len) continue;
        B32 match = 1;
        for (S32 j = 0; j < len; j++) {
            if (toupper((unsigned char)word[j]) != (unsigned char)kw[j]) {
                match = 0;
                break;
            }
        }
        if (match) return sql_keywords[i].type;
    }
    return TOK_IDENTIFIER;
}

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

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

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

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

static Token sql_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 && sql_is_ident_char(*tok->t))
        tok->t++;
    S32 len = (S32)(tok->t - begin);
    token.type = sql_lookup_keyword(begin, len);
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

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

    while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t))
        tok->t++;

    // Decimal part
    if (tok->t < tok->max_t && *tok->t == '.') {
        tok->t++;
        while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t))
            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 sql_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;
}

// Single-quoted string: 'text', with '' as escape for embedded quote
static Token sql_parse_string(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_STRING_LITERAL;

    tok->t++; // skip opening '
    while (tok->t < tok->max_t) {
        if (*tok->t == '\'') {
            tok->t++;
            // '' is an escaped quote inside a string
            if (tok->t < tok->max_t && *tok->t == '\'') {
                tok->t++;
                continue;
            }
            break;
        }
        tok->t++;
    }
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

// Double-quoted identifier: "column_name"
static Token sql_parse_quoted_identifier(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_IDENTIFIER;

    tok->t++; // skip opening "
    while (tok->t < tok->max_t) {
        if (*tok->t == '"') {
            tok->t++;
            // "" is an escaped quote inside a quoted identifier
            if (tok->t < tok->max_t && *tok->t == '"') {
                tok->t++;
                continue;
            }
            break;
        }
        tok->t++;
    }
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

// Backtick-quoted identifier (MySQL): `column_name`
static Token sql_parse_backtick_identifier(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_IDENTIFIER;

    tok->t++; // skip opening `
    while (tok->t < tok->max_t) {
        if (*tok->t == '`') { tok->t++; break; }
        tok->t++;
    }
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

// Line comment: -- ...
static Token sql_parse_line_comment(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_COMMENT;

    tok->t += 2; // skip '--'
    while (tok->t < tok->max_t && *tok->t != '\n') tok->t++;

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

// Block comment: /* ... */
static Token sql_parse_block_comment(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_MULTILINE_COMMENT;

    tok->t += 2; // skip '/*'
    while (tok->t < tok->max_t) {
        if (*tok->t == '*' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '/') {
            tok->t += 2;
            break;
        }
        tok->t++;
    }

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

// MySQL # comment
static Token sql_parse_hash_comment(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_COMMENT;

    tok->t++; // skip '#'
    while (tok->t < tok->max_t && *tok->t != '\n') tok->t++;

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

// Variable / parameter: @var, @@var, :param, $1
static Token sql_parse_variable(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_DIRECTIVE;

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

    if (c == '@') {
        // @@ for system variables
        if (tok->t < tok->max_t && *tok->t == '@') tok->t++;
        while (tok->t < tok->max_t && sql_is_ident_char(*tok->t)) tok->t++;
    } else if (c == ':') {
        while (tok->t < tok->max_t && sql_is_ident_char(*tok->t)) tok->t++;
    } else if (c == '$') {
        // $1, $2, ... (PostgreSQL positional params)
        while (tok->t < tok->max_t && isdigit((unsigned char)*tok->t)) tok->t++;
        // Also handle $tag$ dollar-quoted strings if no digits follow
        if (tok->t == tok->start_t + 1) {
            // Just a lone $, treat as punctuation
            token.type = TOK_PUNCTUATION;
        }
    }

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

static Token sql_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 == ':')) {
            tok->t++;
        }
    }

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

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

static Token sql_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 (sql_is_ident_start(c)) {
        return sql_parse_identifier(tok);
    }
    if (isdigit((unsigned char)c)) {
        return sql_parse_number(tok);
    }

    switch (c) {
    case '\'':
        return sql_parse_string(tok);

    case '"':
        return sql_parse_quoted_identifier(tok);

    case '`':
        return sql_parse_backtick_identifier(tok);

    case '-':
        if ((tok->t + 1) < tok->max_t && *(tok->t + 1) == '-') {
            return sql_parse_line_comment(tok);
        }
        token.type = TOK_OPERATION;
        tok->t++;
        token.len = 1;
        return token;

    case '/':
        if ((tok->t + 1) < tok->max_t && *(tok->t + 1) == '*') {
            return sql_parse_block_comment(tok);
        }
        token.type = TOK_OPERATION;
        tok->t++;
        token.len = 1;
        return token;

    case '#':
        return sql_parse_hash_comment(tok);

    case '.':
        if ((tok->t + 1) < tok->max_t && isdigit((unsigned char)*(tok->t + 1))) {
            return sql_parse_dot_number(tok);
        }
        token.type = TOK_PUNCTUATION;
        tok->t++;
        token.len = 1;
        return token;

    case '@': case '$':
        return sql_parse_variable(tok);

    case ':':
        // :param or :: cast operator
        if ((tok->t + 1) < tok->max_t) {
            char n = *(tok->t + 1);
            if (n == ':') return sql_parse_operator(tok);
            if (sql_is_ident_start(n)) return sql_parse_variable(tok);
        }
        token.type = TOK_PUNCTUATION;
        tok->t++;
        token.len = 1;
        return token;

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

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

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

////////////////////////////////
// tokenize_sql -- main entry point

static void tokenize_sql(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 = sql_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;
    }
}