// lexer_js.c -- JavaScript language tokenizer
//
// Supports nested tagged template literals with ${...} interpolation.
// Uses a depth stack to track template literal nesting so that expressions
// inside ${...} can themselves contain template literals at arbitrary depth.
//
// Example:  html`outer ${css`inner ${x}`} rest`
//           ^^^^ ^^^^^^ ^^^ ^^^^^^ ^ ^  ^^^^^^
//           func string  func str  id str string

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

////////////////////////////////
// JavaScript keyword tables

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

static const JSKeywordEntry js_keywords[] = {
    // Control-flow & language keywords
    {"break",           TOK_KEYWORD},
    {"case",            TOK_KEYWORD},
    {"catch",           TOK_KEYWORD},
    {"class",           TOK_KEYWORD},
    {"const",           TOK_KEYWORD},
    {"continue",        TOK_KEYWORD},
    {"debugger",        TOK_KEYWORD},
    {"default",         TOK_KEYWORD},
    {"delete",          TOK_KEYWORD},
    {"do",              TOK_KEYWORD},
    {"else",            TOK_KEYWORD},
    {"extends",         TOK_KEYWORD},
    {"finally",         TOK_KEYWORD},
    {"for",             TOK_KEYWORD},
    {"function",        TOK_KEYWORD},
    {"if",              TOK_KEYWORD},
    {"in",              TOK_KEYWORD},
    {"instanceof",      TOK_KEYWORD},
    {"let",             TOK_KEYWORD},
    {"new",             TOK_KEYWORD},
    {"of",              TOK_KEYWORD},
    {"return",          TOK_KEYWORD},
    {"switch",          TOK_KEYWORD},
    {"throw",           TOK_KEYWORD},
    {"try",             TOK_KEYWORD},
    {"typeof",          TOK_KEYWORD},
    {"var",             TOK_KEYWORD},
    {"void",            TOK_KEYWORD},
    {"while",           TOK_KEYWORD},
    {"with",            TOK_KEYWORD},
    {"yield",           TOK_KEYWORD},
    {"async",           TOK_KEYWORD},
    {"await",           TOK_KEYWORD},

    // Module keywords (directive-style coloring)
    {"import",          TOK_DIRECTIVE},
    {"export",          TOK_DIRECTIVE},
    {"from",            TOK_DIRECTIVE},
    {"as",              TOK_DIRECTIVE},

    // Values
    {"true",            TOK_VALUE},
    {"false",           TOK_VALUE},
    {"null",            TOK_VALUE},
    {"undefined",       TOK_VALUE},
    {"NaN",             TOK_VALUE},
    {"Infinity",        TOK_VALUE},
    {"this",            TOK_VALUE},
    {"super",           TOK_VALUE},

    // Modifiers / contextual keywords
    {"static",          TOK_MODIFIER},
    {"get",             TOK_MODIFIER},
    {"set",             TOK_MODIFIER},
};

#define JS_KEYWORD_COUNT (S32)(sizeof(js_keywords) / sizeof(js_keywords[0]))

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

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

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

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

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

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

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

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

    char start_char = *tok->t;
    tok->t++;
    if (tok->t >= tok->max_t) goto done;

    if (start_char == '0' && tok->t < tok->max_t) {
        // Hex: 0x / 0X
        if (*tok->t == 'x' || *tok->t == 'X') {
            tok->t++;
            while (tok->t < tok->max_t && (js_is_hex(*tok->t) || *tok->t == '_')) tok->t++;
            if (tok->t < tok->max_t && *tok->t == 'n') tok->t++; // BigInt
            goto done;
        }
        // Octal: 0o / 0O
        if (*tok->t == 'o' || *tok->t == 'O') {
            tok->t++;
            while (tok->t < tok->max_t && ((*tok->t >= '0' && *tok->t <= '7') || *tok->t == '_')) tok->t++;
            if (tok->t < tok->max_t && *tok->t == 'n') tok->t++;
            goto done;
        }
        // Binary: 0b / 0B
        if (*tok->t == 'b' || *tok->t == 'B') {
            tok->t++;
            while (tok->t < tok->max_t && (*tok->t == '0' || *tok->t == '1' || *tok->t == '_')) tok->t++;
            if (tok->t < tok->max_t && *tok->t == 'n') tok->t++;
            goto done;
        }
    }

    // Decimal digits
    while (tok->t < tok->max_t && (isdigit((unsigned char)*tok->t) || *tok->t == '_'))
        tok->t++;

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

    // BigInt suffix
    if (tok->t < tok->max_t && *tok->t == 'n') tok->t++;

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

// Float starting with dot: .5, .123e4
static Token js_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 == '_'))
        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 == '_'))
            tok->t++;
    }

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

// String literal with given quote character (" or ')
static Token js_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;
}

// / followed by / or * (comments), or /= (divide-assign), or bare / (divide)
static Token js_parse_slash_or_comment(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);

    tok->t++; // skip '/'
    if (tok->t >= tok->max_t) {
        token.type = TOK_OPERATION;
        token.len = 1;
        return token;
    }

    if (*tok->t == '/') {
        // Line comment
        token.type = TOK_COMMENT;
        tok->t++;
        while (tok->t < tok->max_t && *tok->t != '\n') tok->t++;
    } else if (*tok->t == '*') {
        // Block comment
        token.type = TOK_MULTILINE_COMMENT;
        tok->t++;
        while (tok->t < tok->max_t) {
            if (*tok->t == '*' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '/') {
                tok->t += 2;
                break;
            }
            tok->t++;
        }
    } else if (*tok->t == '=') {
        // /=
        token.type = TOK_OPERATION;
        tok->t++;
    } else {
        // bare / (division)
        token.type = TOK_OPERATION;
    }

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

// Regex literal: /pattern/flags
static Token js_parse_regex(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_STRING_LITERAL; // color regex like strings

    B32 escape = 0;
    B32 in_class = 0; // inside character class [...]
    tok->t++; // skip opening /

    while (tok->t < tok->max_t && *tok->t != '\n') {
        if (escape) {
            escape = 0;
            tok->t++;
            continue;
        }
        if (*tok->t == '\\') {
            escape = 1;
            tok->t++;
            continue;
        }
        if (*tok->t == '[') { in_class = 1; tok->t++; continue; }
        if (*tok->t == ']') { in_class = 0; tok->t++; continue; }
        if (*tok->t == '/' && !in_class) {
            tok->t++; // closing /
            // Consume flags: d, g, i, m, s, u, v, y
            while (tok->t < tok->max_t && isalpha((unsigned char)*tok->t)) tok->t++;
            break;
        }
        tok->t++;
    }

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

// Operators (handles multi-character sequences)
static Token js_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;

        // Three-character (and four-character) operators
        if ((tok->t + 1) < tok->max_t) {
            char nn = *(tok->t + 1);
            // >>>  and  >>>=
            if (c == '>' && n == '>' && nn == '>') {
                tok->t += 2;
                if (tok->t < tok->max_t && *tok->t == '=') tok->t++;
                goto done;
            }
            // ===
            if (c == '=' && n == '=' && nn == '=') { tok->t += 2; goto done; }
            // !==
            if (c == '!' && n == '=' && nn == '=') { tok->t += 2; goto done; }
            // **=
            if (c == '*' && n == '*' && nn == '=') { tok->t += 2; goto done; }
            // <<=
            if (c == '<' && n == '<' && nn == '=') { tok->t += 2; goto done; }
            // >>=
            if (c == '>' && n == '>' && nn == '=') { tok->t += 2; goto done; }
            // &&=
            if (c == '&' && n == '&' && nn == '=') { tok->t += 2; goto done; }
            // ||=
            if (c == '|' && n == '|' && nn == '=') { tok->t += 2; goto done; }
            // ??=
            if (c == '?' && n == '?' && nn == '=') { tok->t += 2; goto done; }
        }

        // Two-character operators
        if ((c == '=' && n == '=') || (c == '!' && n == '=') ||
            (c == '<' && n == '=') || (c == '>' && n == '=') ||
            (c == '+' && n == '=') || (c == '-' && n == '=') ||
            (c == '*' && n == '=') || (c == '%' && n == '=') ||
            (c == '&' && n == '=') || (c == '|' && n == '=') ||
            (c == '^' && n == '=') || (c == '+' && n == '+') ||
            (c == '-' && n == '-') || (c == '&' && n == '&') ||
            (c == '|' && n == '|') || (c == '<' && n == '<') ||
            (c == '>' && n == '>') || (c == '=' && n == '>') ||
            (c == '*' && n == '*') || (c == '?' && n == '?') ||
            (c == '?' && n == '.')) {
            tok->t++;
        }
    }

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

////////////////////////////////
// Regex vs division heuristic
//
// After identifiers, numbers, values, strings, ) and ] -- slash is division.
// After keywords, operators, most punctuation, start of file -- slash is regex.

static B32 js_slash_is_regex(Token prev, const char *data) {
    switch (prev.type) {
    case TOK_DEFAULT:     return 1; // start of file / no previous token
    case TOK_KEYWORD:     return 1; // e.g. return /regex/
    case TOK_DIRECTIVE:   return 1; // after import/export
    case TOK_OPERATION:
        // After ++ or --, it's division (x++ / y)
        if (prev.len == 2) {
            char c0 = data[prev.start];
            char c1 = data[prev.start + 1];
            if ((c0 == '+' && c1 == '+') || (c0 == '-' && c1 == '-'))
                return 0;
        }
        return 1;
    case TOK_PUNCTUATION:
        // After ) or ], it's division
        if (prev.len == 1) {
            char c = data[prev.start];
            if (c == ')' || c == ']') return 0;
        }
        return 1;
    case TOK_IDENTIFIER:
    case TOK_FUNCTION:
    case TOK_NUMBER:
    case TOK_STRING_LITERAL:
    case TOK_CHAR_LITERAL:
    case TOK_VALUE:
        return 0; // division
    default:
        return 1;
    }
}

////////////////////////////////
// Template literal content scanner
//
// Scans the string content inside a template literal (the text between
// backticks, or between a closing } and the next ${ or closing backtick).
// Paints each byte as TOK_STRING_LITERAL.
//
// Returns 1 if we hit ${  (caller should enter expression mode),
//         0 if we hit closing ` or EOF (template is done).

static int js_scan_template_content(Tokenizer *tok, U8 *out_tokens) {
    while (tok->t < tok->max_t) {
        if (*tok->t == '`') {
            // Closing backtick -- template done
            S32 pos = (S32)(tok->t - tok->buf);
            out_tokens[pos] = TOK_STRING_LITERAL;
            tok->t++;
            return 0;
        }
        if (*tok->t == '$' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '{') {
            // Interpolation start -- paint ${ as punctuation
            S32 pos = (S32)(tok->t - tok->buf);
            out_tokens[pos]     = TOK_PUNCTUATION;
            out_tokens[pos + 1] = TOK_PUNCTUATION;
            tok->t += 2;
            return 1;
        }
        if (*tok->t == '\\' && (tok->t + 1) < tok->max_t) {
            // Escape sequence -- paint both chars as string
            S32 pos = (S32)(tok->t - tok->buf);
            out_tokens[pos]     = TOK_STRING_LITERAL;
            out_tokens[pos + 1] = TOK_STRING_LITERAL;
            tok->t += 2;
            continue;
        }
        // Regular string character
        S32 pos = (S32)(tok->t - tok->buf);
        out_tokens[pos] = TOK_STRING_LITERAL;
        tok->t++;
    }
    return 0; // EOF -- unclosed template
}

////////////////////////////////
// Get next non-template token
//
// Called by the main loop for normal expression parsing.  Template backticks
// and interpolation-closing braces are handled directly by the main loop
// before this function is reached.

static Token js_get_next_token(Tokenizer *tok, Token prev, const char *data) {
    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 (js_is_ident_start(c)) {
        return js_parse_identifier(tok);
    }
    if (isdigit((unsigned char)c)) {
        return js_parse_number(tok);
    }

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

    case '/':
        // Check for comment first (// or /*)
        if ((tok->t + 1) < tok->max_t) {
            char n = *(tok->t + 1);
            if (n == '/' || n == '*')
                return js_parse_slash_or_comment(tok);
        }
        // Regex or division based on context
        if (js_slash_is_regex(prev, data))
            return js_parse_regex(tok);
        return js_parse_slash_or_comment(tok);

    // Dot: float literal (.5), spread (...), or member access
    case '.':
        if ((tok->t + 1) < tok->max_t && isdigit((unsigned char)*(tok->t + 1))) {
            return js_parse_dot_number(tok);
        }
        if ((tok->t + 2) < tok->max_t && *(tok->t + 1) == '.' && *(tok->t + 2) == '.') {
            token.type = TOK_OPERATION;
            tok->t += 3;
            token.len = 3;
            return token;
        }
        token.type = TOK_PUNCTUATION;
        tok->t++;
        token.len = 1;
        return token;

    // Punctuation
    case ';': 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 '^': case '~':
    case '?':
        return js_parse_operator(tok);

    // Private class fields: #name
    case '#':
        token.type = TOK_IDENTIFIER;
        tok->t++;
        while (tok->t < tok->max_t && js_is_ident_char(*tok->t)) tok->t++;
        token.len = (S32)(tok->t - tok->start_t);
        return token;

    // Decorators: @name
    case '@':
        token.type = TOK_DIRECTIVE;
        tok->t++;
        while (tok->t < tok->max_t && js_is_ident_char(*tok->t)) tok->t++;
        token.len = (S32)(tok->t - tok->start_t);
        return token;

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

////////////////////////////////
// tokenize_js -- main entry point
//
// Tokenizes the full buffer and paints out_tokens[] with token types.
//
// Template literal nesting is tracked with a depth stack.  Each level
// records how many unmatched { braces exist in the current interpolation
// expression.  When a } is encountered and the brace count is zero, it
// closes the interpolation and we resume scanning template string content.
//
//   html`text ${obj.x} more ${css`inner ${y}`} end`
//        |str  ||expr | str  ||    |str  ||  || str|
//               ^punc  ^punc  ^punc      ^   ^^
//               ${     }      ${         }   }` (nesting!)

#define JS_MAX_TEMPLATE_DEPTH 32

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

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

    S32 tmpl_depth = 0;
    S32 brace_count[JS_MAX_TEMPLATE_DEPTH];
    memset(brace_count, 0, sizeof(brace_count));

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

    while (tok.t < tok.max_t) {
        tokenizer_eat_whitespace(&tok);
        if (tok.t >= tok.max_t) break;

        char c = *tok.t;

        // ---- Template literal: opening backtick ----
        // This handles both top-level template literals and nested ones
        // (e.g. a tagged template inside a ${...} interpolation).
        if (c == '`') {
            S32 pos = (S32)(tok.t - tok.buf);
            out_tokens[pos] = TOK_STRING_LITERAL;
            tok.t++;

            int result = js_scan_template_content(&tok, out_tokens);
            if (result == 1) {
                // Hit ${  -- push template depth
                if (tmpl_depth < JS_MAX_TEMPLATE_DEPTH) {
                    brace_count[tmpl_depth] = 0;
                    tmpl_depth++;
                }
            }
            // result == 0: self-contained template (no interpolation, or
            // all interpolations already resolved recursively)

            prev.type = TOK_STRING_LITERAL;
            prev.start = pos;
            prev.len = (S32)(tok.t - tok.buf) - pos;
            continue;
        }

        // ---- Closing brace that ends a template interpolation ----
        // When we're inside a template expression and the brace count is
        // zero, this } closes the ${...} and we resume string scanning.
        if (c == '}' && tmpl_depth > 0 && brace_count[tmpl_depth - 1] == 0) {
            S32 pos = (S32)(tok.t - tok.buf);
            out_tokens[pos] = TOK_PUNCTUATION;
            tok.t++;
            tmpl_depth--;

            // Resume template string content
            int result = js_scan_template_content(&tok, out_tokens);
            if (result == 1) {
                // Hit another ${  -- re-enter expression mode
                if (tmpl_depth < JS_MAX_TEMPLATE_DEPTH) {
                    brace_count[tmpl_depth] = 0;
                    tmpl_depth++;
                }
            }
            // result == 0: template closed with backtick

            prev.type = TOK_STRING_LITERAL;
            prev.start = pos;
            prev.len = (S32)(tok.t - tok.buf) - pos;
            continue;
        }

        // ---- Normal token ----
        tok.start_t = tok.t;
        Token token = js_get_next_token(&tok, prev, data);
        if (token.len == 0) break;

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

        // Track brace depth for template interpolation
        if (tmpl_depth > 0 && token.type == TOK_PUNCTUATION && token.len == 1) {
            char tc = data[token.start];
            if (tc == '{')      brace_count[tmpl_depth - 1]++;
            else if (tc == '}') brace_count[tmpl_depth - 1]--;
        }

        // Retroactively mark identifier before '(' as 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;
    }
}