// lexer_go.c -- Go language tokenizer
//
// Same pattern as lexer_c.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>

////////////////////////////////
// Go keyword tables

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

static const GoKeywordEntry go_keywords[] = {
    // Control-flow & language keywords
    {"break",           TOK_KEYWORD},
    {"case",            TOK_KEYWORD},
    {"chan",             TOK_KEYWORD},
    {"const",           TOK_KEYWORD},
    {"continue",        TOK_KEYWORD},
    {"default",         TOK_KEYWORD},
    {"defer",           TOK_KEYWORD},
    {"else",            TOK_KEYWORD},
    {"fallthrough",     TOK_KEYWORD},
    {"for",             TOK_KEYWORD},
    {"func",            TOK_KEYWORD},
    {"go",              TOK_KEYWORD},
    {"goto",            TOK_KEYWORD},
    {"if",              TOK_KEYWORD},
    {"import",          TOK_KEYWORD},
    {"interface",       TOK_KEYWORD},
    {"map",             TOK_KEYWORD},
    {"package",         TOK_KEYWORD},
    {"range",           TOK_KEYWORD},
    {"return",          TOK_KEYWORD},
    {"select",          TOK_KEYWORD},
    {"struct",          TOK_KEYWORD},
    {"switch",          TOK_KEYWORD},
    {"type",            TOK_KEYWORD},
    {"var",             TOK_KEYWORD},

    // Built-in types
    {"bool",            TOK_TYPE},
    {"byte",            TOK_TYPE},
    {"complex64",       TOK_TYPE},
    {"complex128",      TOK_TYPE},
    {"error",           TOK_TYPE},
    {"float32",         TOK_TYPE},
    {"float64",         TOK_TYPE},
    {"int",             TOK_TYPE},
    {"int8",            TOK_TYPE},
    {"int16",           TOK_TYPE},
    {"int32",           TOK_TYPE},
    {"int64",           TOK_TYPE},
    {"rune",            TOK_TYPE},
    {"string",          TOK_TYPE},
    {"uint",            TOK_TYPE},
    {"uint8",           TOK_TYPE},
    {"uint16",          TOK_TYPE},
    {"uint32",          TOK_TYPE},
    {"uint64",          TOK_TYPE},
    {"uintptr",         TOK_TYPE},
    {"any",             TOK_TYPE},
    {"comparable",      TOK_TYPE},

    // Built-in values
    {"true",            TOK_VALUE},
    {"false",           TOK_VALUE},
    {"nil",             TOK_VALUE},
    {"iota",            TOK_VALUE},

    // Built-in functions (treated as modifiers for distinct coloring)
    {"append",          TOK_MODIFIER},
    {"cap",             TOK_MODIFIER},
    {"clear",           TOK_MODIFIER},
    {"close",           TOK_MODIFIER},
    {"complex",         TOK_MODIFIER},
    {"copy",            TOK_MODIFIER},
    {"delete",          TOK_MODIFIER},
    {"imag",            TOK_MODIFIER},
    {"len",             TOK_MODIFIER},
    {"make",            TOK_MODIFIER},
    {"max",             TOK_MODIFIER},
    {"min",             TOK_MODIFIER},
    {"new",             TOK_MODIFIER},
    {"panic",           TOK_MODIFIER},
    {"print",           TOK_MODIFIER},
    {"println",         TOK_MODIFIER},
    {"real",            TOK_MODIFIER},
    {"recover",         TOK_MODIFIER},
};

#define GO_KEYWORD_COUNT (S32)(sizeof(go_keywords) / sizeof(go_keywords[0]))

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

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

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

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

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

static B32 go_is_octal(char c) {
    return c >= '0' && c <= '7';
}

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

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

static Token go_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 or 0X
        if (*tok->t == 'x' || *tok->t == 'X') {
            tok->t++;
            while (tok->t < tok->max_t && (go_is_hex(*tok->t) || *tok->t == '_')) tok->t++;
            // Hex float: optional .hex_digits, optional p exponent
            if (tok->t < tok->max_t && *tok->t == '.') {
                tok->t++;
                while (tok->t < tok->max_t && (go_is_hex(*tok->t) || *tok->t == '_')) tok->t++;
            }
            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 == '_')) tok->t++;
            }
            goto imaginary;
        }
        // Octal: 0o or 0O
        if (*tok->t == 'o' || *tok->t == 'O') {
            tok->t++;
            while (tok->t < tok->max_t && (go_is_octal(*tok->t) || *tok->t == '_')) tok->t++;
            goto imaginary;
        }
        // Binary: 0b or 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++;
            goto imaginary;
        }
    }

    // Decimal or float (also handles legacy octal like 0755)
    while (tok->t < tok->max_t && (isdigit((unsigned char)*tok->t) || *tok->t == '_'))
        tok->t++;

    // Decimal float: .digits
    if (tok->t < tok->max_t && *tok->t == '.') {
        // Check next char is digit to avoid consuming '..' or method calls
        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 == '_')) tok->t++;
        } else if ((tok->t + 1) >= tok->max_t || !go_is_ident_start(*(tok->t + 1))) {
            // Trailing dot like "1." -- still a float in Go
            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++;
    }

imaginary:
    // Imaginary suffix
    if (tok->t < tok->max_t && *tok->t == 'i')
        tok->t++;

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

// Dot-starting float: .5, .123e4
static Token go_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++;
    // 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++;
    }
    // Imaginary
    if (tok->t < tok->max_t && *tok->t == 'i')
        tok->t++;

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

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

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

// Raw string literal: `...`
static Token go_parse_raw_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++; break; }
        tok->t++;
    }
    token.len = (S32)(tok->t - tok->start_t);
    return token;
}

// Rune literal: 'x', '\n', '\u0041', etc.
static Token go_parse_rune_literal(Tokenizer *tok) {
    Token token;
    token.start = (S32)(tok->start_t - tok->buf);
    token.type = TOK_CHAR_LITERAL;

    tok->t++; // skip opening '
    if (tok->t < tok->max_t && *tok->t == '\\') {
        tok->t++; // skip backslash
        // Consume escape sequence characters
        if (tok->t < tok->max_t) {
            char esc = *tok->t;
            tok->t++;
            if (esc == 'x') {
                // \xNN
                for (int i = 0; i < 2 && tok->t < tok->max_t && go_is_hex(*tok->t); i++) tok->t++;
            } else if (esc == 'u') {
                // \uNNNN
                for (int i = 0; i < 4 && tok->t < tok->max_t && go_is_hex(*tok->t); i++) tok->t++;
            } else if (esc == 'U') {
                // \UNNNNNNNN
                for (int i = 0; i < 8 && tok->t < tok->max_t && go_is_hex(*tok->t); i++) tok->t++;
            } else if (go_is_octal(esc)) {
                // \NNN
                for (int i = 0; i < 2 && tok->t < tok->max_t && go_is_octal(*tok->t); i++) tok->t++;
            }
            // else: simple escape like \n, \t, \\, \' -- already consumed
        }
    } else if (tok->t < tok->max_t && *tok->t != '\n' && *tok->t != '\'') {
        tok->t++; // the rune character
    }
    if (tok->t < tok->max_t && *tok->t == '\'') tok->t++; // closing '

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

static Token go_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 {
        token.type = TOK_OPERATION;
    }

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

static Token go_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 operators
        if (c == '<' && n == '<' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '=') {
            tok->t += 2; goto done;
        }
        if (c == '>' && n == '>' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '=') {
            tok->t += 2; goto done;
        }
        if (c == '&' && n == '^' && (tok->t + 1) < tok->max_t && *(tok->t + 1) == '=') {
            tok->t += 2; goto done;
        }
        // := (short variable declaration)
        if (c == ':' && n == '=') { tok->t++; 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 == '^')) {
            tok->t++;
        }
    }

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

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

static Token go_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; // EOF
    }

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

    if (go_is_ident_start(c)) {
        return go_parse_identifier(tok);
    }
    if (isdigit((unsigned char)c)) {
        return go_parse_number(tok);
    }

    switch (c) {
    case '"':  return go_parse_string(tok);
    case '`':  return go_parse_raw_string(tok);
    case '\'': return go_parse_rune_literal(tok);
    case '/':  return go_parse_slash_or_comment(tok);

    // Dot: could start a float literal like .5
    case '.':
        if ((tok->t + 1) < tok->max_t && isdigit((unsigned char)*(tok->t + 1))) {
            return go_parse_dot_number(tok);
        }
        // Ellipsis ... or just punctuation
        token.type = TOK_PUNCTUATION;
        tok->t++;
        if (tok->t < tok->max_t && *tok->t == '.' &&
            (tok->t + 1) < tok->max_t && *(tok->t + 1) == '.') {
            tok->t += 2; // consume ...
        }
        token.len = (S32)(tok->t - tok->start_t);
        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 '^': case '~':
    case ':':
        return go_parse_operator(tok);

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

////////////////////////////////
// tokenize_go -- main entry point
//
// Tokenizes the full buffer and paints out_tokens[] with token types.
// Second pass: retroactively marks identifiers before '(' as functions.

static void tokenize_go(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 = go_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;
    }
}