/*
 * par2serial-cc: lexer.c - Lexical Analyzer Implementation
 */
#include "lexer.h"

/* ── Keyword table ───────────────────────────────────────── */
typedef struct { const char *kw; TokenType type; } KWEntry;

static const KWEntry keywords[] = {
    /* C keywords */
    {"auto",     TOK_AUTO},     {"break",    TOK_BREAK},
    {"case",     TOK_CASE},     {"char",     TOK_CHAR},
    {"const",    TOK_CONST},    {"continue", TOK_CONTINUE},
    {"default",  TOK_DEFAULT},  {"do",       TOK_DO},
    {"double",   TOK_DOUBLE},   {"else",     TOK_ELSE},
    {"enum",     TOK_ENUM},     {"extern",   TOK_EXTERN},
    {"float",    TOK_FLOAT},    {"for",      TOK_FOR},
    {"goto",     TOK_GOTO},     {"if",       TOK_IF},
    {"inline",   TOK_INLINE},   {"int",      TOK_INT},
    {"long",     TOK_LONG},     {"register", TOK_REGISTER},
    {"restrict", TOK_RESTRICT}, {"return",   TOK_RETURN},
    {"short",    TOK_SHORT},    {"signed",   TOK_SIGNED},
    {"sizeof",   TOK_SIZEOF},   {"static",   TOK_STATIC},
    {"struct",   TOK_STRUCT},   {"switch",   TOK_SWITCH},
    {"typedef",  TOK_TYPEDEF},  {"union",    TOK_UNION},
    {"unsigned", TOK_UNSIGNED}, {"void",     TOK_VOID},
    {"volatile", TOK_VOLATILE}, {"while",    TOK_WHILE},

    /* Parallel extensions */
    {"parallel_for",    TOK_PARALLEL_FOR},
    {"parallel_reduce", TOK_PARALLEL_REDUCE},
    {"parallel_scan",   TOK_PARALLEL_SCAN},
    {"parallel_map",    TOK_PARALLEL_MAP},
    {"barrier",         TOK_BARRIER},
    {"shared",          TOK_SHARED},
    {"tile_hint",       TOK_TILE_HINT},
    {"simd_hint",       TOK_SIMD_HINT},
    {"memory_layout",   TOK_MEMORY_LAYOUT},
    {"atomic_add",      TOK_ATOMIC_ADD},
    {"atomic_cas",      TOK_ATOMIC_CAS},
    {"thread_id",       TOK_THREAD_ID},
    {"num_threads",     TOK_NUM_THREADS},
    {NULL, TOK_EOF}
};

/* ── Helper macros ───────────────────────────────────────── */
#define PEEK(lex)       ((lex)->pos < (lex)->src_len ? (lex)->src[(lex)->pos] : '\0')
#define PEEK_AT(lex, n) ((lex)->pos + (n) < (lex)->src_len ? (lex)->src[(lex)->pos + (n)] : '\0')
#define AT_END(lex)     ((lex)->pos >= (lex)->src_len)

static char advance(Lexer *lex) {
    char c = PEEK(lex);
    if (c == '\n') { lex->line++; lex->col = 1; }
    else { lex->col++; }
    lex->pos++;
    return c;
}

static bool match(Lexer *lex, char expected) {
    if (PEEK(lex) == expected) { advance(lex); return true; }
    return false;
}

static SourceLoc make_loc(Lexer *lex) {
    return (SourceLoc){lex->filename, lex->line, lex->col};
}

/* ── Token creation ──────────────────────────────────────── */

static Token make_token(Lexer *lex, TokenType type, const char *start, size_t len, SourceLoc loc) {
    Token t;
    t.type = type;
    t.loc  = loc;
    t.text = arena_strndup(lex->arena, start, len);
    t.text_len = len;
    t.int_val = 0;
    return t;
}

static void emit(Lexer *lex, Token t) {
    vec_push(&lex->tokens, t);
}

/* ── Skip whitespace and comments ────────────────────────── */
static void skip_ws(Lexer *lex) {
    while (!AT_END(lex)) {
        char c = PEEK(lex);
        if (c == ' ' || c == '\t' || c == '\r' || c == '\n') {
            advance(lex);
        } else if (c == '/' && PEEK_AT(lex, 1) == '/') {
            /* line comment */
            while (!AT_END(lex) && PEEK(lex) != '\n') advance(lex);
        } else if (c == '/' && PEEK_AT(lex, 1) == '*') {
            /* block comment */
            advance(lex); advance(lex);
            while (!AT_END(lex)) {
                if (PEEK(lex) == '*' && PEEK_AT(lex, 1) == '/') {
                    advance(lex); advance(lex);
                    break;
                }
                advance(lex);
            }
        } else {
            break;
        }
    }
}

/* ── Lex number ──────────────────────────────────────────── */
static void lex_number(Lexer *lex) {
    SourceLoc loc = make_loc(lex);
    const char *start = lex->src + lex->pos;
    bool is_float = false;
    bool is_hex = false;

    if (PEEK(lex) == '0' && (PEEK_AT(lex, 1) == 'x' || PEEK_AT(lex, 1) == 'X')) {
        advance(lex); advance(lex);
        is_hex = true;
        while (isxdigit(PEEK(lex))) advance(lex);
    } else {
        while (isdigit(PEEK(lex))) advance(lex);
        if (PEEK(lex) == '.' && isdigit(PEEK_AT(lex, 1))) {
            is_float = true;
            advance(lex);
            while (isdigit(PEEK(lex))) advance(lex);
        }
        if (PEEK(lex) == 'e' || PEEK(lex) == 'E') {
            is_float = true;
            advance(lex);
            if (PEEK(lex) == '+' || PEEK(lex) == '-') advance(lex);
            while (isdigit(PEEK(lex))) advance(lex);
        }
    }
    /* suffix: f, l, u, ll etc */
    while (PEEK(lex) == 'f' || PEEK(lex) == 'F' ||
           PEEK(lex) == 'l' || PEEK(lex) == 'L' ||
           PEEK(lex) == 'u' || PEEK(lex) == 'U')
        advance(lex);

    size_t len = (lex->src + lex->pos) - start;
    Token t = make_token(lex, is_float ? TOK_FLOAT_LIT : TOK_INT_LIT, start, len, loc);
    if (is_float) {
        t.float_val = strtod(start, NULL);
    } else if (is_hex) {
        t.int_val = (int64_t)strtoll(start, NULL, 16);
    } else {
        t.int_val = (int64_t)strtoll(start, NULL, 10);
    }
    emit(lex, t);
}

/* ── Lex string ──────────────────────────────────────────── */
static void lex_string(Lexer *lex) {
    SourceLoc loc = make_loc(lex);
    advance(lex); /* skip opening " */
    const char *start = lex->src + lex->pos;
    while (!AT_END(lex) && PEEK(lex) != '"') {
        if (PEEK(lex) == '\\') advance(lex); /* skip escaped char */
        advance(lex);
    }
    size_t len = (lex->src + lex->pos) - start;
    if (!AT_END(lex)) advance(lex); /* skip closing " */
    emit(lex, make_token(lex, TOK_STRING_LIT, start, len, loc));
}

/* ── Lex char literal ────────────────────────────────────── */
static void lex_char(Lexer *lex) {
    SourceLoc loc = make_loc(lex);
    advance(lex); /* skip ' */
    const char *start = lex->src + lex->pos;
    if (PEEK(lex) == '\\') advance(lex);
    advance(lex);
    size_t len = (lex->src + lex->pos) - start;
    if (!AT_END(lex)) advance(lex); /* skip closing ' */
    Token t = make_token(lex, TOK_CHAR_LIT, start, len, loc);
    t.int_val = (int64_t)(unsigned char)start[0];
    emit(lex, t);
}

/* ── Lex identifier/keyword ──────────────────────────────── */
static void lex_ident(Lexer *lex) {
    SourceLoc loc = make_loc(lex);
    const char *start = lex->src + lex->pos;
    while (isalnum(PEEK(lex)) || PEEK(lex) == '_') advance(lex);
    size_t len = (lex->src + lex->pos) - start;

    /* check keywords */
    TokenType type = TOK_IDENT;
    for (const KWEntry *k = keywords; k->kw; k++) {
        if (strlen(k->kw) == len && strncmp(k->kw, start, len) == 0) {
            type = k->type;
            break;
        }
    }
    emit(lex, make_token(lex, type, start, len, loc));
}

/* ── Lex preprocessor line ───────────────────────────────── */
static void lex_preprocessor(Lexer *lex) {
    SourceLoc loc = make_loc(lex);
    const char *start = lex->src + lex->pos;
    advance(lex); /* skip # */
    while (!AT_END(lex) && PEEK(lex) != '\n') {
        if (PEEK(lex) == '\\' && PEEK_AT(lex, 1) == '\n') {
            advance(lex); advance(lex); /* line continuation */
        } else {
            advance(lex);
        }
    }
    size_t len = (lex->src + lex->pos) - start;
    emit(lex, make_token(lex, TOK_HASH, start, len, loc));
}

/* ── Main tokenize ───────────────────────────────────────── */

Lexer *lexer_create(Arena *arena, const char *src, size_t len, const char *filename) {
    Lexer *lex = (Lexer *)arena_alloc(arena, sizeof(Lexer));
    lex->src = src;
    lex->src_len = len;
    lex->pos = 0;
    lex->line = 1;
    lex->col = 1;
    lex->filename = filename;
    lex->arena = arena;
    vec_init(&lex->tokens);
    lex->error_count = 0;
    return lex;
}

bool lexer_tokenize(Lexer *lex) {
    while (!AT_END(lex)) {
        skip_ws(lex);
        if (AT_END(lex)) break;

        char c = PEEK(lex);
        SourceLoc loc = make_loc(lex);

        /* Numbers */
        if (isdigit(c) || (c == '.' && isdigit(PEEK_AT(lex, 1)))) {
            lex_number(lex);
            continue;
        }

        /* Identifiers / keywords */
        if (isalpha(c) || c == '_') {
            lex_ident(lex);
            continue;
        }

        /* String literal */
        if (c == '"') { lex_string(lex); continue; }

        /* Char literal */
        if (c == '\'') { lex_char(lex); continue; }

        /* Preprocessor */
        if (c == '#') { lex_preprocessor(lex); continue; }

        /* Multi-char operators and single-char tokens */
        advance(lex);
        switch (c) {
        case '(':  emit(lex, make_token(lex, TOK_LPAREN,   &c, 1, loc)); break;
        case ')':  emit(lex, make_token(lex, TOK_RPAREN,   &c, 1, loc)); break;
        case '{':  emit(lex, make_token(lex, TOK_LBRACE,   &c, 1, loc)); break;
        case '}':  emit(lex, make_token(lex, TOK_RBRACE,   &c, 1, loc)); break;
        case '[':  emit(lex, make_token(lex, TOK_LBRACKET, &c, 1, loc)); break;
        case ']':  emit(lex, make_token(lex, TOK_RBRACKET, &c, 1, loc)); break;
        case ';':  emit(lex, make_token(lex, TOK_SEMICOLON,&c, 1, loc)); break;
        case ',':  emit(lex, make_token(lex, TOK_COMMA,    &c, 1, loc)); break;
        case '~':  emit(lex, make_token(lex, TOK_TILDE,    &c, 1, loc)); break;
        case '?':  emit(lex, make_token(lex, TOK_QUESTION, &c, 1, loc)); break;
        case ':':  emit(lex, make_token(lex, TOK_COLON,    &c, 1, loc)); break;

        case '+':
            if (match(lex, '+'))      emit(lex, make_token(lex, TOK_INC,     "++", 2, loc));
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_PLUS_EQ, "+=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_PLUS,    "+",  1, loc));
            break;
        case '-':
            if (match(lex, '-'))      emit(lex, make_token(lex, TOK_DEC,      "--", 2, loc));
            else if (match(lex, '>')) emit(lex, make_token(lex, TOK_ARROW,    "->", 2, loc));
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_MINUS_EQ, "-=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_MINUS,    "-",  1, loc));
            break;
        case '*':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_STAR_EQ, "*=", 2, loc));
            else                 emit(lex, make_token(lex, TOK_STAR,    "*",  1, loc));
            break;
        case '/':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_SLASH_EQ, "/=", 2, loc));
            else                 emit(lex, make_token(lex, TOK_SLASH,    "/",  1, loc));
            break;
        case '%':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_PERCENT_EQ, "%=", 2, loc));
            else                 emit(lex, make_token(lex, TOK_PERCENT,    "%",  1, loc));
            break;
        case '&':
            if (match(lex, '&'))      emit(lex, make_token(lex, TOK_AND,    "&&", 2, loc));
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_AMP_EQ, "&=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_AMP,    "&",  1, loc));
            break;
        case '|':
            if (match(lex, '|'))      emit(lex, make_token(lex, TOK_OR,      "||", 2, loc));
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_PIPE_EQ, "|=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_PIPE,    "|",  1, loc));
            break;
        case '^':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_CARET_EQ, "^=", 2, loc));
            else                 emit(lex, make_token(lex, TOK_CARET,    "^",  1, loc));
            break;
        case '=':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_EQ,     "==", 2, loc));
            else                 emit(lex, make_token(lex, TOK_ASSIGN, "=",  1, loc));
            break;
        case '!':
            if (match(lex, '=')) emit(lex, make_token(lex, TOK_NEQ,  "!=", 2, loc));
            else                 emit(lex, make_token(lex, TOK_BANG, "!",  1, loc));
            break;
        case '<':
            if (match(lex, '<')) {
                if (match(lex, '=')) emit(lex, make_token(lex, TOK_LSHIFT_EQ, "<<=", 3, loc));
                else                 emit(lex, make_token(lex, TOK_LSHIFT,    "<<",  2, loc));
            }
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_LE, "<=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_LT, "<",  1, loc));
            break;
        case '>':
            if (match(lex, '>')) {
                if (match(lex, '=')) emit(lex, make_token(lex, TOK_RSHIFT_EQ, ">>=", 3, loc));
                else                 emit(lex, make_token(lex, TOK_RSHIFT,    ">>",  2, loc));
            }
            else if (match(lex, '=')) emit(lex, make_token(lex, TOK_GE, ">=", 2, loc));
            else                      emit(lex, make_token(lex, TOK_GT, ">",  1, loc));
            break;
        case '.':
            if (PEEK(lex) == '.' && PEEK_AT(lex, 1) == '.') {
                advance(lex); advance(lex);
                emit(lex, make_token(lex, TOK_ELLIPSIS, "...", 3, loc));
            } else {
                emit(lex, make_token(lex, TOK_DOT, ".", 1, loc));
            }
            break;

        default:
            p2s_error(loc, "unexpected character '%c' (0x%02X)", c, (unsigned char)c);
            lex->error_count++;
            break;
        }
    }

    /* Append EOF */
    SourceLoc eof_loc = make_loc(lex);
    emit(lex, make_token(lex, TOK_EOF, "", 0, eof_loc));
    return lex->error_count == 0;
}

/* ── Debug ───────────────────────────────────────────────── */
static const char *token_names[] = {
    [TOK_INT_LIT]="INT_LIT", [TOK_FLOAT_LIT]="FLOAT_LIT",
    [TOK_STRING_LIT]="STRING", [TOK_CHAR_LIT]="CHAR",
    [TOK_IDENT]="IDENT",
    [TOK_AUTO]="auto", [TOK_BREAK]="break", [TOK_CASE]="case",
    [TOK_CHAR]="char", [TOK_CONST]="const", [TOK_CONTINUE]="continue",
    [TOK_DEFAULT]="default", [TOK_DO]="do", [TOK_DOUBLE]="double",
    [TOK_ELSE]="else", [TOK_ENUM]="enum", [TOK_EXTERN]="extern",
    [TOK_FLOAT]="float", [TOK_FOR]="for", [TOK_GOTO]="goto",
    [TOK_IF]="if", [TOK_INLINE]="inline", [TOK_INT]="int",
    [TOK_LONG]="long", [TOK_REGISTER]="register", [TOK_RESTRICT]="restrict",
    [TOK_RETURN]="return", [TOK_SHORT]="short", [TOK_SIGNED]="signed",
    [TOK_SIZEOF]="sizeof", [TOK_STATIC]="static", [TOK_STRUCT]="struct",
    [TOK_SWITCH]="switch", [TOK_TYPEDEF]="typedef", [TOK_UNION]="union",
    [TOK_UNSIGNED]="unsigned", [TOK_VOID]="void", [TOK_VOLATILE]="volatile",
    [TOK_WHILE]="while",
    [TOK_PARALLEL_FOR]="parallel_for", [TOK_PARALLEL_REDUCE]="parallel_reduce",
    [TOK_PARALLEL_SCAN]="parallel_scan", [TOK_PARALLEL_MAP]="parallel_map",
    [TOK_BARRIER]="barrier", [TOK_SHARED]="shared",
    [TOK_TILE_HINT]="tile_hint", [TOK_SIMD_HINT]="simd_hint",
    [TOK_MEMORY_LAYOUT]="memory_layout",
    [TOK_ATOMIC_ADD]="atomic_add", [TOK_ATOMIC_CAS]="atomic_cas",
    [TOK_THREAD_ID]="thread_id", [TOK_NUM_THREADS]="num_threads",
    [TOK_PLUS]="+", [TOK_MINUS]="-", [TOK_STAR]="*", [TOK_SLASH]="/",
    [TOK_PERCENT]="%", [TOK_AMP]="&", [TOK_PIPE]="|", [TOK_CARET]="^",
    [TOK_TILDE]="~", [TOK_BANG]="!", [TOK_LT]="<", [TOK_GT]=">",
    [TOK_EQ]="==", [TOK_NEQ]="!=", [TOK_LE]="<=", [TOK_GE]=">=",
    [TOK_AND]="&&", [TOK_OR]="||", [TOK_LSHIFT]="<<", [TOK_RSHIFT]=">>",
    [TOK_ASSIGN]="=",
    [TOK_PLUS_EQ]="+=", [TOK_MINUS_EQ]="-=", [TOK_STAR_EQ]="*=",
    [TOK_SLASH_EQ]="/=", [TOK_PERCENT_EQ]="%=", [TOK_AMP_EQ]="&=",
    [TOK_PIPE_EQ]="|=", [TOK_CARET_EQ]="^=",
    [TOK_LSHIFT_EQ]="<<=", [TOK_RSHIFT_EQ]=">>=",
    [TOK_INC]="++", [TOK_DEC]="--", [TOK_ARROW]="->", [TOK_DOT]=".",
    [TOK_QUESTION]="?", [TOK_COLON]=":",
    [TOK_LPAREN]="(", [TOK_RPAREN]=")", [TOK_LBRACE]="{", [TOK_RBRACE]="}",
    [TOK_LBRACKET]="[", [TOK_RBRACKET]="]", [TOK_SEMICOLON]=";",
    [TOK_COMMA]=",", [TOK_ELLIPSIS]="...", [TOK_HASH]="#",
    [TOK_EOF]="EOF", [TOK_ERROR]="ERROR",
};

const char *token_type_str(TokenType t) {
    if (t >= 0 && t < TOK_COUNT && token_names[t])
        return token_names[t];
    return "UNKNOWN";
}

void token_print(const Token *t) {
    printf("%-18s %-20.*s  [%d:%d]\n",
           token_type_str(t->type),
           (int)t->text_len, t->text,
           t->loc.line, t->loc.col);
}