/*
** PostgreSQL SQL Tokenizer for Lime Parser
**
** Hand-written tokenizer that implements the lexical analysis described in
** PostgreSQL's scan.l, adapted for use with a Lime-generated parser.
**
** This tokenizer handles:
**   - SQL keywords (case-insensitive lookup via binary search)
**   - Identifiers (lowercased, truncated at PG_NAMEDATALEN)
**   - Integer, float, and numeric literals (decimal, hex, octal, binary)
**   - String literals (standard, extended E'', dollar-quoted, bit, hex, Unicode)
**   - C-style nested block comments and SQL line comments
**   - All PostgreSQL operator and punctuation tokens
**   - String continuation across newlines (SQL standard)
**   - Location tracking (byte offset)
**   - One-token lookahead for LALR(1) fixups (FORMAT_LA, NOT_LA, etc.)
**
** See LEXER_NOTES.md for the detailed analysis this is based on.
*/

#define _POSIX_C_SOURCE 200809L

#include "tokenize.h"
#include "pg_kwlist.h"         /* Auto-generated keyword table (uses pg_grammar.h) */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <errno.h>
#include <limits.h>

/* Maximum identifier length (same as PostgreSQL NAMEDATALEN) */
#define PG_NAMEDATALEN 64

/* Initial and growth factor for the literal buffer */
#define LITBUF_INIT 256
#define LITBUF_GROW 2


/* ======================================================================
** Keyword table
**
** The keyword table is provided by pg_kwlist.h (auto-generated by
** gen_kwlist.py from pg_grammar.h + gram.lime).  It contains 510
** entries sorted alphabetically, each with a Lime token ID, reservation
** category, and bare-label flag.
**
** The pg_kw_lookup() inline function from pg_kwlist.h performs binary
** search.  We wrap it below to match the public API declared in
** tokenize.h.
** ======================================================================
*/

/* ======================================================================
** Single-character token mapping
**
** Lime uses named tokens for all single characters that PostgreSQL's
** lexer returns as their ASCII value.  This function maps the character
** to the corresponding Lime token ID from pg_grammar.h.
** ====================================================================== */

static int self_to_token(unsigned char c) {
    switch (c) {
        case '(': return LPAREN;
        case ')': return RPAREN;
        case '[': return LBRACKET;
        case ']': return RBRACKET;
        case ',': return COMMA;
        case ';': return SEMICOLON;
        case ':': return COLON;
        case '.': return DOT;
        case '+': return PLUS;
        case '-': return MINUS;
        case '*': return STAR;
        case '/': return SLASH;
        case '%': return PERCENT;
        case '^': return CARET;
        case '<': return LT;
        case '>': return GT;
        case '=': return EQ;
        case '|': return PIPE;
        case '#': return HASH;
        default:  return self_to_token(c);  /* fallback for unexpected chars */
    }
}


/* Keyword lookup / name functions using pg_kwlist.h */

/* The keyword table is now in pg_kwlist.h (auto-generated). */

/*
** Case-insensitive keyword comparison for binary search.
** The pg_kwlist.h table is already lowercase; we lowercase input on the fly.
*/
static int kw_cmp_ci(const char *input, size_t len, const char *keyword) {
    for (size_t i = 0; i < len; i++) {
        unsigned char a = (unsigned char)input[i];
        unsigned char b = (unsigned char)keyword[i];
        if (b == '\0') return 1;   /* input longer than keyword */
        if (a >= 'A' && a <= 'Z') a += 32;
        if (a != b) return (int)a - (int)b;
    }
    if (keyword[len] != '\0') return -1;
    return 0;
}

/*
** Look up a keyword by name (case-insensitive).
** Returns the Lime token code, or 0 if not a keyword.
*/
int pg_keyword_lookup(const char *text, size_t len) {
    int lo = 0;
    int hi = PG_NUM_KEYWORDS - 1;
    while (lo <= hi) {
        int mid = lo + (hi - lo) / 2;
        int cmp = kw_cmp_ci(text, len, pg_keywords[mid].keyword);
        if (cmp < 0)
            hi = mid - 1;
        else if (cmp > 0)
            lo = mid + 1;
        else
            return pg_keywords[mid].token_id;
    }
    return 0;  /* not a keyword */
}

/*
** Get the canonical (lowercase) keyword name for a token code.
** Linear scan since token codes are not necessarily sequential.
*/
const char *pg_keyword_name(int token) {
    for (int i = 0; i < PG_NUM_KEYWORDS; i++) {
        if (pg_keywords[i].token_id == token)
            return pg_keywords[i].keyword;
    }
    return NULL;
}

/*
** Return the number of keywords in the table.
*/
int pg_keyword_count(void) {
    return PG_NUM_KEYWORDS;
}


/* ======================================================================
** Scanner state
** ====================================================================== */

struct PgScanner {
    const char *input;      /* Input SQL string */
    size_t      length;     /* Length of input */
    size_t      pos;        /* Current byte offset */

    /* Literal buffer for string accumulation */
    char       *litbuf;
    size_t      litlen;
    size_t      litalloc;

    /* Lookahead buffer for pg_scan_with_lookahead() */
    bool        have_lookahead;     /* True if lookahead token is buffered */
    int         la_token;           /* Buffered lookahead token code */
    PgTokenValue la_val;            /* Buffered lookahead semantic value */
    int         la_loc;             /* Buffered lookahead location */

    /* Error state */
    const char *errmsg;
    int         errloc;     /* Byte offset of error, or -1 */
};


/* ======================================================================
** Literal buffer helpers
** ====================================================================== */

static void litbuf_init(PgScanner *s) {
    s->litlen = 0;
}

static bool litbuf_ensure(PgScanner *s, size_t need) {
    size_t required = s->litlen + need;
    if (required < s->litalloc) return true;
    size_t newsize = s->litalloc;
    while (newsize <= required)
        newsize *= LITBUF_GROW;
    char *newbuf = realloc(s->litbuf, newsize);
    if (!newbuf) return false;
    s->litbuf = newbuf;
    s->litalloc = newsize;
    return true;
}

static void litbuf_add(PgScanner *s, const char *text, size_t len)
    __attribute__((unused));
static void litbuf_add(PgScanner *s, const char *text, size_t len) {
    if (!litbuf_ensure(s, len)) return;
    memcpy(s->litbuf + s->litlen, text, len);
    s->litlen += len;
}

static void litbuf_addchar(PgScanner *s, char c) {
    if (!litbuf_ensure(s, 1)) return;
    s->litbuf[s->litlen++] = c;
}

static char *litbuf_dup(PgScanner *s) {
    char *result = malloc(s->litlen + 1);
    if (!result) return NULL;
    memcpy(result, s->litbuf, s->litlen);
    result[s->litlen] = '\0';
    return result;
}


/* ======================================================================
** Character classification helpers
** ====================================================================== */

static inline bool is_space(unsigned char c) {
    return c == ' ' || c == '\t' || c == '\n' || c == '\r' ||
           c == '\f' || c == '\v';
}

static inline bool is_newline(unsigned char c) {
    return c == '\n' || c == '\r';
}

static inline bool is_ident_start(unsigned char c) {
    return (c >= 'A' && c <= 'Z') ||
           (c >= 'a' && c <= 'z') ||
           c == '_' || c >= 0x80;
}

static inline bool is_ident_cont(unsigned char c) {
    return (c >= 'A' && c <= 'Z') ||
           (c >= 'a' && c <= 'z') ||
           (c >= '0' && c <= '9') ||
           c == '_' || c == '$' || c >= 0x80;
}

static inline bool is_digit(unsigned char c) {
    return c >= '0' && c <= '9';
}

static inline bool is_hexdigit(unsigned char c) {
    return (c >= '0' && c <= '9') ||
           (c >= 'a' && c <= 'f') ||
           (c >= 'A' && c <= 'F');
}

static inline bool is_octdigit(unsigned char c) {
    return c >= '0' && c <= '7';
}

static inline bool is_bindigit(unsigned char c) {
    return c == '0' || c == '1';
}

/* Characters that are returned as single-character tokens */
static inline bool is_self(unsigned char c) {
    return c == ',' || c == '(' || c == ')' || c == '[' || c == ']' ||
           c == '.' || c == ';' || c == ':' || c == '|' || c == '+' ||
           c == '-' || c == '*' || c == '/' || c == '%' || c == '^' ||
           c == '<' || c == '>' || c == '=';
}

/* Characters that can appear in operators */
static inline bool is_op_char(unsigned char c) {
    return c == '~' || c == '!' || c == '@' || c == '#' || c == '^' ||
           c == '&' || c == '|' || c == '`' || c == '?' || c == '+' ||
           c == '-' || c == '*' || c == '/' || c == '%' || c == '<' ||
           c == '>' || c == '=';
}

/* Characters that make an operator "non-SQL-standard" (allowing trailing +/-) */
static inline bool is_special_op_char(unsigned char c) {
    return c == '~' || c == '!' || c == '@' || c == '#' || c == '^' ||
           c == '&' || c == '|' || c == '`' || c == '?' || c == '%';
}

static inline unsigned char to_lower(unsigned char c) {
    if (c >= 'A' && c <= 'Z') return c + 32;
    return c;
}


/* ======================================================================
** Forward declarations
** ====================================================================== */

static int scan_string_body(PgScanner *s, int start_loc);
static int scan_extended_string_body(PgScanner *s, int start_loc);
static int scan_bitstring_body(PgScanner *s, char prefix, int start_loc);
static int scan_delimited_ident(PgScanner *s, int start_loc);
static int scan_unicode_ident(PgScanner *s, int start_loc);
static int scan_unicode_string(PgScanner *s, int start_loc);
static void skip_block_comment(PgScanner *s);


/* ======================================================================
** Skip whitespace and comments
**
** Returns true if any whitespace containing a newline was seen.
** This is needed for the string continuation check.
** ====================================================================== */

static bool skip_whitespace(PgScanner *s) {
    bool saw_newline = false;
    while (s->pos < s->length) {
        unsigned char c = (unsigned char)s->input[s->pos];

        if (is_space(c)) {
            if (is_newline(c)) saw_newline = true;
            s->pos++;
            continue;
        }

        /* SQL line comment: -- to end of line */
        if (c == '-' && s->pos + 1 < s->length && s->input[s->pos + 1] == '-') {
            s->pos += 2;
            while (s->pos < s->length && !is_newline((unsigned char)s->input[s->pos]))
                s->pos++;
            /* The newline itself is consumed on the next iteration */
            continue;
        }

        /* C-style block comment: / * ... * / with nesting */
        if (c == '/' && s->pos + 1 < s->length && s->input[s->pos + 1] == '*') {
            s->pos += 2;
            skip_block_comment(s);
            continue;
        }

        break;
    }
    return saw_newline;
}

static void skip_block_comment(PgScanner *s) {
    int depth = 1;
    while (s->pos < s->length && depth > 0) {
        unsigned char c = (unsigned char)s->input[s->pos];
        if (c == '/' && s->pos + 1 < s->length && s->input[s->pos + 1] == '*') {
            depth++;
            s->pos += 2;
        } else if (c == '*' && s->pos + 1 < s->length && s->input[s->pos + 1] == '/') {
            depth--;
            s->pos += 2;
        } else {
            s->pos++;
        }
    }
    if (depth > 0) {
        s->errmsg = "unterminated /* comment";
        s->errloc = (int)s->pos;
    }
}


/* ======================================================================
** Check for string continuation (SQL standard)
**
** After a closing quote, if whitespace with at least one newline
** is followed by another single quote, the string continues.
** ====================================================================== */

static bool check_string_continuation(PgScanner *s) {
    size_t saved = s->pos;
    bool saw_newline = false;

    /* Skip whitespace and comments, tracking newlines */
    while (s->pos < s->length) {
        unsigned char c = (unsigned char)s->input[s->pos];
        if (is_space(c)) {
            if (is_newline(c)) saw_newline = true;
            s->pos++;
            continue;
        }
        if (c == '-' && s->pos + 1 < s->length && s->input[s->pos + 1] == '-') {
            s->pos += 2;
            while (s->pos < s->length && !is_newline((unsigned char)s->input[s->pos]))
                s->pos++;
            continue;
        }
        if (c == '/' && s->pos + 1 < s->length && s->input[s->pos + 1] == '*') {
            s->pos += 2;
            skip_block_comment(s);
            continue;
        }
        break;
    }

    if (saw_newline && s->pos < s->length && s->input[s->pos] == '\'') {
        s->pos++;  /* consume the opening quote of continuation */
        return true;
    }

    /* No continuation -- restore position */
    s->pos = saved;
    return false;
}


/* ======================================================================
** String literal scanning
** ====================================================================== */

/*
** Scan a standard single-quoted string body (already past opening quote).
** Handles '' escaping and string continuation.
** Returns SCONST on success, 0 on error.
*/
static int scan_string_body(PgScanner *s, int start_loc) {
    (void)start_loc;
    litbuf_init(s);

    for (;;) {
        while (s->pos < s->length) {
            char c = s->input[s->pos];
            if (c == '\'') {
                s->pos++;
                /* Check for '' (embedded quote) */
                if (s->pos < s->length && s->input[s->pos] == '\'') {
                    litbuf_addchar(s, '\'');
                    s->pos++;
                    continue;
                }
                /* End of string body -- check for continuation */
                if (check_string_continuation(s)) {
                    /* Continue accumulating into the same literal */
                    continue;
                }
                return SCONST;
            }
            litbuf_addchar(s, c);
            s->pos++;
        }
        /* Reached EOF inside string */
        s->errmsg = "unterminated quoted string";
        s->errloc = start_loc;
        return 0;
    }
}

/*
** Scan an extended (E'...') string body with backslash escapes.
*/
static int scan_extended_string_body(PgScanner *s, int start_loc) {
    litbuf_init(s);

    for (;;) {
        while (s->pos < s->length) {
            unsigned char c = (unsigned char)s->input[s->pos];

            if (c == '\'') {
                s->pos++;
                /* Check for '' */
                if (s->pos < s->length && s->input[s->pos] == '\'') {
                    litbuf_addchar(s, '\'');
                    s->pos++;
                    continue;
                }
                /* End of string -- check continuation */
                if (check_string_continuation(s)) {
                    continue;
                }
                return SCONST;
            }

            if (c == '\\') {
                s->pos++;
                if (s->pos >= s->length) {
                    litbuf_addchar(s, '\\');
                    break;
                }
                unsigned char esc = (unsigned char)s->input[s->pos];
                switch (esc) {
                    case 'b':  litbuf_addchar(s, '\b'); s->pos++; break;
                    case 'f':  litbuf_addchar(s, '\f'); s->pos++; break;
                    case 'n':  litbuf_addchar(s, '\n'); s->pos++; break;
                    case 'r':  litbuf_addchar(s, '\r'); s->pos++; break;
                    case 't':  litbuf_addchar(s, '\t'); s->pos++; break;
                    case 'v':  litbuf_addchar(s, '\v'); s->pos++; break;
                    case '\\': litbuf_addchar(s, '\\'); s->pos++; break;
                    case '\'': litbuf_addchar(s, '\''); s->pos++; break;

                    /* Octal escape: \NNN */
                    case '0': case '1': case '2': case '3':
                    case '4': case '5': case '6': case '7': {
                        unsigned int val = esc - '0';
                        s->pos++;
                        if (s->pos < s->length && is_octdigit((unsigned char)s->input[s->pos])) {
                            val = val * 8 + ((unsigned char)s->input[s->pos] - '0');
                            s->pos++;
                            if (s->pos < s->length && is_octdigit((unsigned char)s->input[s->pos])) {
                                val = val * 8 + ((unsigned char)s->input[s->pos] - '0');
                                s->pos++;
                            }
                        }
                        litbuf_addchar(s, (char)(unsigned char)val);
                        break;
                    }

                    /* Hex escape: \xNN */
                    case 'x': {
                        s->pos++;
                        if (s->pos < s->length && is_hexdigit((unsigned char)s->input[s->pos])) {
                            unsigned int val = 0;
                            unsigned char h = (unsigned char)s->input[s->pos];
                            if (h >= '0' && h <= '9') val = h - '0';
                            else if (h >= 'a' && h <= 'f') val = h - 'a' + 10;
                            else val = h - 'A' + 10;
                            s->pos++;
                            if (s->pos < s->length && is_hexdigit((unsigned char)s->input[s->pos])) {
                                h = (unsigned char)s->input[s->pos];
                                val *= 16;
                                if (h >= '0' && h <= '9') val += h - '0';
                                else if (h >= 'a' && h <= 'f') val += h - 'a' + 10;
                                else val += h - 'A' + 10;
                                s->pos++;
                            }
                            litbuf_addchar(s, (char)(unsigned char)val);
                        } else {
                            litbuf_addchar(s, 'x');
                        }
                        break;
                    }

                    /* Unicode escapes: \uXXXX or \UXXXXXXXX */
                    case 'u': case 'U': {
                        int ndigits = (esc == 'u') ? 4 : 8;
                        s->pos++;
                        /* For this initial implementation, just pass the
                        ** Unicode escape through as-is.  Full conversion
                        ** to UTF-8 can be added later. */
                        litbuf_addchar(s, '\\');
                        litbuf_addchar(s, (char)esc);
                        for (int i = 0; i < ndigits && s->pos < s->length; i++) {
                            if (!is_hexdigit((unsigned char)s->input[s->pos])) break;
                            litbuf_addchar(s, s->input[s->pos]);
                            s->pos++;
                        }
                        break;
                    }

                    default:
                        litbuf_addchar(s, (char)esc);
                        s->pos++;
                        break;
                }
                continue;
            }

            litbuf_addchar(s, (char)c);
            s->pos++;
        }
        s->errmsg = "unterminated quoted string";
        s->errloc = start_loc;
        return 0;
    }
}

/* Dollar-quoted string scanning is handled inline in pg_scan() */

/*
** Scan bit/hex string body (B'...' or X'...').
** prefix is 'b' or 'x'.
*/
static int scan_bitstring_body(PgScanner *s, char prefix, int start_loc) {
    litbuf_init(s);
    litbuf_addchar(s, prefix);

    while (s->pos < s->length) {
        char c = s->input[s->pos];
        if (c == '\'') {
            s->pos++;
            /* Check for '' */
            if (s->pos < s->length && s->input[s->pos] == '\'') {
                litbuf_addchar(s, '\'');
                s->pos++;
                continue;
            }
            /* Check for continuation */
            if (check_string_continuation(s)) {
                continue;
            }
            return (prefix == 'b') ? BCONST : XCONST;
        }
        litbuf_addchar(s, c);
        s->pos++;
    }
    s->errmsg = (prefix == 'b')
        ? "unterminated bit string literal"
        : "unterminated hexadecimal string literal";
    s->errloc = start_loc;
    return 0;
}

/*
** Scan a double-quoted (delimited) identifier.
** Called after the opening " has been consumed.
*/
static int scan_delimited_ident(PgScanner *s, int start_loc) {
    litbuf_init(s);

    while (s->pos < s->length) {
        char c = s->input[s->pos];
        if (c == '"') {
            s->pos++;
            /* Check for "" */
            if (s->pos < s->length && s->input[s->pos] == '"') {
                litbuf_addchar(s, '"');
                s->pos++;
                continue;
            }
            if (s->litlen == 0) {
                s->errmsg = "zero-length delimited identifier";
                s->errloc = start_loc;
                return 0;
            }
            return IDENT;
        }
        litbuf_addchar(s, c);
        s->pos++;
    }
    s->errmsg = "unterminated quoted identifier";
    s->errloc = start_loc;
    return 0;
}

/*
** Scan a Unicode-escape identifier: U&"..."
** Called after 'U&"' has been consumed.
*/
static int scan_unicode_ident(PgScanner *s, int start_loc) {
    litbuf_init(s);

    while (s->pos < s->length) {
        char c = s->input[s->pos];
        if (c == '"') {
            s->pos++;
            if (s->pos < s->length && s->input[s->pos] == '"') {
                litbuf_addchar(s, '"');
                s->pos++;
                continue;
            }
            if (s->litlen == 0) {
                s->errmsg = "zero-length delimited identifier";
                s->errloc = start_loc;
                return 0;
            }
            return UIDENT;
        }
        litbuf_addchar(s, c);
        s->pos++;
    }
    s->errmsg = "unterminated quoted identifier";
    s->errloc = start_loc;
    return 0;
}

/*
** Scan a Unicode-escape string: U&'...'
** Called after U&' has been consumed.
*/
static int scan_unicode_string(PgScanner *s, int start_loc) {
    litbuf_init(s);

    while (s->pos < s->length) {
        char c = s->input[s->pos];
        if (c == '\'') {
            s->pos++;
            if (s->pos < s->length && s->input[s->pos] == '\'') {
                litbuf_addchar(s, '\'');
                s->pos++;
                continue;
            }
            if (check_string_continuation(s)) {
                continue;
            }
            return USCONST;
        }
        litbuf_addchar(s, c);
        s->pos++;
    }
    s->errmsg = "unterminated quoted string";
    s->errloc = start_loc;
    return 0;
}


/* ======================================================================
** Number scanning helpers
** ====================================================================== */

/*
** Scan decimal digits with optional underscore separators.
** Returns number of digits consumed (not counting underscores).
*/
static int scan_digits(PgScanner *s, bool (*is_valid)(unsigned char)) {
    int count = 0;
    while (s->pos < s->length) {
        unsigned char c = (unsigned char)s->input[s->pos];
        if (is_valid(c)) {
            count++;
            s->pos++;
        } else if (c == '_' && count > 0) {
            /* Underscore must be followed by a valid digit */
            if (s->pos + 1 < s->length && is_valid((unsigned char)s->input[s->pos + 1])) {
                s->pos++;  /* skip underscore */
            } else {
                break;
            }
        } else {
            break;
        }
    }
    return count;
}

/*
** Try to parse the accumulated number text as an int32.
** Returns true if it fits, false if overflow.
*/
static bool try_parse_int32(const char *text, size_t len, int base, int32_t *out) {
    /* Copy to NUL-terminated buffer, stripping underscores */
    char buf[128];
    size_t j = 0;
    size_t start = 0;

    /* Skip base prefix */
    if (base != 10 && len >= 2) {
        start = 2;  /* skip 0x, 0o, 0b */
    }

    for (size_t i = start; i < len && j < sizeof(buf) - 1; i++) {
        if (text[i] != '_')
            buf[j++] = text[i];
    }
    buf[j] = '\0';

    errno = 0;
    char *end;
    long val = strtol(buf, &end, base);
    if (errno == ERANGE || val < INT32_MIN || val > INT32_MAX)
        return false;
    if (*end != '\0')
        return false;
    *out = (int32_t)val;
    return true;
}


/* ======================================================================
** Main scanner function
**
** pg_scan() extracts the next token from the input.
**
** Returns:
**   > 0: Token code
**   0:   End of input
**   -1:  Error (check s->errmsg)
**
** On success, fills in *val with the semantic value and *loc with
** the byte offset of the token start.
** ====================================================================== */

int pg_scan(PgScanner *s, PgTokenValue *val, int *loc) {
    unsigned char c;

restart:
    skip_whitespace(s);

    if (s->pos >= s->length) {
        *loc = (int)s->pos;
        return 0;  /* EOF */
    }

    *loc = (int)s->pos;
    c = (unsigned char)s->input[s->pos];

    /* ---- Identifiers and keywords ---- */
    if (is_ident_start(c)) {
        size_t start = s->pos;
        s->pos++;
        while (s->pos < s->length && is_ident_cont((unsigned char)s->input[s->pos]))
            s->pos++;
        size_t len = s->pos - start;

        /* Check for string-type prefixes: B'', X'', E'', N'', U&'', U&"" */
        if (len == 1 && s->pos < s->length && s->input[s->pos] == '\'') {
            unsigned char first = to_lower((unsigned char)s->input[start]);
            if (first == 'b') {
                s->pos++;  /* consume quote */
                int tok = scan_bitstring_body(s, 'b', (int)start);
                if (tok == 0) return -1;
                val->str = litbuf_dup(s);
                return tok;
            }
            if (first == 'x') {
                s->pos++;
                int tok = scan_bitstring_body(s, 'x', (int)start);
                if (tok == 0) return -1;
                val->str = litbuf_dup(s);
                return tok;
            }
            if (first == 'e') {
                s->pos++;
                int tok = scan_extended_string_body(s, (int)start);
                if (tok == 0) return -1;
                val->str = litbuf_dup(s);
                return tok;
            }
            if (first == 'n') {
                /* National character: treat as keyword NCHAR + push back quote */
                s->pos = start + 1;  /* back to just after 'n' */
                int kwtoken = pg_keyword_lookup("nchar", 5);
                if (kwtoken > 0) {
                    val->keyword = pg_keyword_name(kwtoken);
                    return kwtoken;
                }
                /* NCHAR not a keyword -- return 'n' as identifier */
                val->str = strdup("n");
                return IDENT;
            }
        }

        /* Check for U& prefix (Unicode escapes) */
        if (len == 1 && to_lower((unsigned char)s->input[start]) == 'u' &&
            s->pos < s->length && s->input[s->pos] == '&') {
            if (s->pos + 1 < s->length) {
                char next = s->input[s->pos + 1];
                if (next == '\'') {
                    s->pos += 2;  /* skip &' */
                    int tok = scan_unicode_string(s, (int)start);
                    if (tok == 0) return -1;
                    val->str = litbuf_dup(s);
                    return tok;
                }
                if (next == '"') {
                    s->pos += 2;  /* skip &" */
                    int tok = scan_unicode_ident(s, (int)start);
                    if (tok == 0) return -1;
                    val->str = litbuf_dup(s);
                    return tok;
                }
            }
            /* Not U& followed by quote -- fall through to identifier */
        }

        /* Keyword lookup */
        int kwtoken = pg_keyword_lookup(s->input + start, len);
        if (kwtoken > 0) {
            val->keyword = pg_keyword_name(kwtoken);
            return kwtoken;
        }

        /* Regular identifier: lowercase and truncate */
        size_t ident_len = (len < PG_NAMEDATALEN - 1) ? len : PG_NAMEDATALEN - 1;
        char *ident = malloc(ident_len + 1);
        if (!ident) return -1;
        for (size_t i = 0; i < ident_len; i++)
            ident[i] = (char)to_lower((unsigned char)s->input[start + i]);
        ident[ident_len] = '\0';
        val->str = ident;
        return IDENT;
    }

    /* ---- Numeric literals ---- */
    if (is_digit(c) || (c == '.' && s->pos + 1 < s->length &&
                        is_digit((unsigned char)s->input[s->pos + 1]))) {
        size_t start = s->pos;
        bool is_float = false;
        int base = 10;

        if (c == '0' && s->pos + 1 < s->length) {
            unsigned char next = to_lower((unsigned char)s->input[s->pos + 1]);
            if (next == 'x') {
                base = 16; s->pos += 2;
                if (s->pos < s->length && s->input[s->pos] == '_') s->pos++;
                if (scan_digits(s, is_hexdigit) == 0) {
                    s->errmsg = "invalid hexadecimal integer";
                    s->errloc = (int)start;
                    return -1;
                }
                goto check_integer_junk;
            }
            if (next == 'o') {
                base = 8; s->pos += 2;
                if (s->pos < s->length && s->input[s->pos] == '_') s->pos++;
                if (scan_digits(s, is_octdigit) == 0) {
                    s->errmsg = "invalid octal integer";
                    s->errloc = (int)start;
                    return -1;
                }
                goto check_integer_junk;
            }
            if (next == 'b' && s->pos + 2 < s->length &&
                is_bindigit((unsigned char)s->input[s->pos + 2])) {
                /* Only treat as binary if followed by binary digit
                ** (to distinguish from 0b being start of identifier) */
                base = 2; s->pos += 2;
                if (s->pos < s->length && s->input[s->pos] == '_') s->pos++;
                if (scan_digits(s, is_bindigit) == 0) {
                    s->errmsg = "invalid binary integer";
                    s->errloc = (int)start;
                    return -1;
                }
                goto check_integer_junk;
            }
        }

        /* Decimal integer or start of numeric/real */
        if (c != '.') {
            scan_digits(s, is_digit);
        }

        /* Check for decimal point */
        if (s->pos < s->length && s->input[s->pos] == '.') {
            /* Check for '..' (DOT_DOT) - don't consume the dots */
            if (s->pos + 1 < s->length && s->input[s->pos + 1] == '.') {
                /* This is numericfail: return the integer part, leave '..' */
                goto finish_integer;
            }
            is_float = true;
            s->pos++;  /* consume '.' */
            scan_digits(s, is_digit);
        } else if (c == '.') {
            is_float = true;
            s->pos++;  /* consume leading '.' */
            scan_digits(s, is_digit);
        }

        /* Check for exponent */
        if (s->pos < s->length &&
            (s->input[s->pos] == 'e' || s->input[s->pos] == 'E')) {
            size_t exp_start = s->pos;
            s->pos++;
            if (s->pos < s->length &&
                (s->input[s->pos] == '+' || s->input[s->pos] == '-')) {
                s->pos++;
            }
            if (s->pos >= s->length || !is_digit((unsigned char)s->input[s->pos])) {
                /* realfail: no digits after exponent */
                if (is_float) {
                    s->errmsg = "trailing junk after numeric literal";
                    s->errloc = (int)exp_start;
                    return -1;
                }
                /* Revert: not an exponent */
                s->pos = exp_start;
                goto finish_integer;
            }
            is_float = true;
            scan_digits(s, is_digit);
        }

        if (is_float) {
            /* Check for trailing identifier chars (junk) */
            if (s->pos < s->length && is_ident_start((unsigned char)s->input[s->pos])) {
                s->errmsg = "trailing junk after numeric literal";
                s->errloc = (int)s->pos;
                return -1;
            }
            size_t len = s->pos - start;
            val->str = strndup(s->input + start, len);
            return FCONST;
        }

    finish_integer:
        ;  /* empty statement after label */

    check_integer_junk:
        /* Check for trailing identifier chars */
        if (s->pos < s->length && is_ident_start((unsigned char)s->input[s->pos])) {
            s->errmsg = "trailing junk after numeric literal";
            s->errloc = (int)s->pos;
            return -1;
        }

        size_t numlen = s->pos - start;
        int32_t ival;
        if (try_parse_int32(s->input + start, numlen, base, &ival)) {
            val->ival = ival;
            return ICONST;
        }
        /* Overflow: treat as float */
        val->str = strndup(s->input + start, numlen);
        return FCONST;
    }

    /* ---- String literals ---- */
    if (c == '\'') {
        s->pos++;
        int tok = scan_string_body(s, (int)(s->pos - 1));
        if (tok == 0) return -1;
        val->str = litbuf_dup(s);
        return tok;
    }

    /* ---- Dollar-quoted strings ---- */
    if (c == '$') {
        /* Check for $tag$ pattern */
        size_t dstart = s->pos;
        s->pos++;  /* past initial $ */

        /* Read optional tag */
        size_t tag_start = s->pos;
        if (s->pos < s->length && is_ident_start((unsigned char)s->input[s->pos])) {
            s->pos++;
            while (s->pos < s->length && is_ident_cont((unsigned char)s->input[s->pos]) &&
                   s->input[s->pos] != '$')
                s->pos++;
        }

        if (s->pos < s->length && s->input[s->pos] == '$') {
            /* We have a valid dollar-quote delimiter */
            size_t delim_len = s->pos - dstart + 1;
            s->pos++;  /* past closing $ of delimiter */

            /* Save the delimiter for matching */
            char *delim = strndup(s->input + dstart, delim_len);
            if (!delim) return -1;

            /* Scan for matching delimiter */
            litbuf_init(s);
            while (s->pos < s->length) {
                if (s->input[s->pos] == '$') {
                    /* Check if this starts the closing delimiter */
                    if (s->pos + delim_len <= s->length &&
                        memcmp(s->input + s->pos, delim, delim_len) == 0) {
                        s->pos += delim_len;
                        free(delim);
                        val->str = litbuf_dup(s);
                        return SCONST;
                    }
                }
                litbuf_addchar(s, s->input[s->pos]);
                s->pos++;
            }
            free(delim);
            s->errmsg = "unterminated dollar-quoted string";
            s->errloc = (int)dstart;
            return -1;
        }

        /* Not a dollar-quote delimiter.  Check for $N parameter. */
        s->pos = dstart + 1;  /* back to after the $ */
        if (s->pos < s->length && is_digit((unsigned char)s->input[s->pos])) {
            size_t pstart = s->pos;
            while (s->pos < s->length && is_digit((unsigned char)s->input[s->pos]))
                s->pos++;
            /* Check for junk */
            if (s->pos < s->length && is_ident_start((unsigned char)s->input[s->pos])) {
                s->errmsg = "trailing junk after parameter";
                s->errloc = (int)dstart;
                return -1;
            }
            char *pstr = strndup(s->input + pstart, s->pos - pstart);
            if (!pstr) return -1;
            errno = 0;
            long pnum = strtol(pstr, NULL, 10);
            free(pstr);
            if (errno == ERANGE || pnum > INT32_MAX || pnum < 0) {
                s->errmsg = "parameter number too large";
                s->errloc = (int)dstart;
                return -1;
            }
            val->ival = (int32_t)pnum;
            return PARAM;
        }

        /* Plain $ -- return as single character token.
        ** The grammar does not use bare $, so this will be a syntax error. */
        s->pos = dstart + 1;
        (void)tag_start;
        return self_to_token('$');
    }

    /* ---- Double-quoted identifier ---- */
    if (c == '"') {
        s->pos++;
        int tok = scan_delimited_ident(s, (int)(s->pos - 1));
        if (tok == 0) return -1;
        val->str = litbuf_dup(s);
        return tok;
    }

    /* ---- Multi-character operators (check before single-char) ---- */

    /* :: (TYPECAST) */
    if (c == ':' && s->pos + 1 < s->length && s->input[s->pos + 1] == ':') {
        s->pos += 2;
        return TYPECAST;
    }

    /* := (COLON_EQUALS) */
    if (c == ':' && s->pos + 1 < s->length && s->input[s->pos + 1] == '=') {
        s->pos += 2;
        return COLON_EQUALS;
    }

    /* .. (DOT_DOT) */
    if (c == '.' && s->pos + 1 < s->length && s->input[s->pos + 1] == '.') {
        s->pos += 2;
        return DOT_DOT;
    }

    /* => (EQUALS_GREATER) */
    if (c == '=' && s->pos + 1 < s->length && s->input[s->pos + 1] == '>') {
        s->pos += 2;
        return EQUALS_GREATER;
    }

    /* >= (GREATER_EQUALS) */
    if (c == '>' && s->pos + 1 < s->length && s->input[s->pos + 1] == '=') {
        s->pos += 2;
        return GREATER_EQUALS;
    }

    /* <= (LESS_EQUALS) */
    if (c == '<' && s->pos + 1 < s->length && s->input[s->pos + 1] == '=') {
        s->pos += 2;
        return LESS_EQUALS;
    }

    /* <> (NOT_EQUALS) */
    if (c == '<' && s->pos + 1 < s->length && s->input[s->pos + 1] == '>') {
        s->pos += 2;
        return NOT_EQUALS;
    }

    /* != (NOT_EQUALS) */
    if (c == '!' && s->pos + 1 < s->length && s->input[s->pos + 1] == '=') {
        s->pos += 2;
        return NOT_EQUALS;
    }

    /* -> (RIGHT_ARROW) */
    if (c == '-' && s->pos + 1 < s->length && s->input[s->pos + 1] == '>') {
        s->pos += 2;
        return RIGHT_ARROW;
    }

    /* ---- Generic operator ---- */
    if (is_op_char(c) && !is_self(c)) {
        /* This character is an op_char but not a self char, so it must
        ** start an operator (e.g. ~, !, @, #, &, `, ?) */
        goto scan_operator;
    }

    /* For characters that are in BOTH op_chars and self, handle them
    ** as operators only if followed by another op_char. */
    if (is_op_char(c) && is_self(c)) {
        if (s->pos + 1 < s->length && is_op_char((unsigned char)s->input[s->pos + 1])) {
            /* But first check for comment starts */
            if (c == '-' && s->input[s->pos + 1] == '-') {
                /* Line comment -- should have been caught by skip_whitespace.
                ** This shouldn't happen, but handle gracefully. */
                goto restart;
            }
            if (c == '/' && s->input[s->pos + 1] == '*') {
                goto restart;
            }
            goto scan_operator;
        }
        /* Single self character */
        s->pos++;
        return self_to_token(c);
    }

    if (is_self(c)) {
        s->pos++;
        return self_to_token(c);
    }

    /* ---- Scan a multi-character operator ---- */
    if (false) {
    scan_operator:
        ;
        size_t op_start = s->pos;
        s->pos++;
        while (s->pos < s->length && is_op_char((unsigned char)s->input[s->pos])) {
            /* Check for embedded comment starts */
            if (s->input[s->pos] == '/' && s->pos + 1 < s->length &&
                s->input[s->pos + 1] == '*') break;
            if (s->input[s->pos] == '-' && s->pos + 1 < s->length &&
                s->input[s->pos + 1] == '-') break;
            s->pos++;
        }

        size_t nchars = s->pos - op_start;

        /* SQL compatibility: + and - cannot end a multi-char operator
        ** unless it contains a "special" op char */
        if (nchars > 1) {
            char last = s->input[s->pos - 1];
            if (last == '+' || last == '-') {
                bool has_special = false;
                for (size_t i = op_start; i < s->pos - 1; i++) {
                    if (is_special_op_char((unsigned char)s->input[i])) {
                        has_special = true;
                        break;
                    }
                }
                if (!has_special) {
                    /* Strip trailing + and - */
                    while (nchars > 1) {
                        char ch = s->input[op_start + nchars - 1];
                        if (ch != '+' && ch != '-') break;
                        nchars--;
                    }
                    s->pos = op_start + nchars;
                }
            }
        }

        /* If reduced to single self char, return it as character token */
        if (nchars == 1 && is_self((unsigned char)s->input[op_start])) {
            return self_to_token((unsigned char)s->input[op_start]);
        }

        /* If reduced to a known 2-char token, return that */
        if (nchars == 2) {
            char c0 = s->input[op_start], c1 = s->input[op_start + 1];
            if (c0 == '=' && c1 == '>') return EQUALS_GREATER;
            if (c0 == '>' && c1 == '=') return GREATER_EQUALS;
            if (c0 == '<' && c1 == '=') return LESS_EQUALS;
            if (c0 == '<' && c1 == '>') return NOT_EQUALS;
            if (c0 == '!' && c1 == '=') return NOT_EQUALS;
            if (c0 == '-' && c1 == '>') return RIGHT_ARROW;
        }

        if (nchars >= PG_NAMEDATALEN) {
            s->errmsg = "operator too long";
            s->errloc = (int)op_start;
            return -1;
        }

        val->str = strndup(s->input + op_start, nchars);
        return OP;
    }

    /* ---- Unknown character ---- */
    s->pos++;
    return self_to_token(c);
}


/* ======================================================================
** Public API
** ====================================================================== */

/*
** Create and initialize a scanner.
*/
PgScanner *pg_scanner_create(const char *input, size_t length) {
    PgScanner *s = calloc(1, sizeof(PgScanner));
    if (!s) return NULL;
    s->input = input;
    s->length = length;
    s->pos = 0;
    s->litalloc = LITBUF_INIT;
    s->litbuf = malloc(LITBUF_INIT);
    if (!s->litbuf) { free(s); return NULL; }
    s->litlen = 0;
    s->have_lookahead = false;
    s->errmsg = NULL;
    s->errloc = -1;
    return s;
}

/*
** Destroy a scanner and free resources.
*/
void pg_scanner_destroy(PgScanner *s) {
    if (!s) return;
    free(s->litbuf);
    free(s);
}

/*
** Get the last error message, or NULL if no error.
*/
const char *pg_scanner_error(const PgScanner *s) {
    return s->errmsg;
}

/*
** Get the byte offset of the last error, or -1.
*/
int pg_scanner_error_location(const PgScanner *s) {
    return s->errloc;
}

/*
** Get token type name for debugging.
*/
const char *pg_token_name(int token) {
    switch (token) {
        case 0:               return "EOF";
        case ICONST:       return "ICONST";
        case FCONST:       return "FCONST";
        case SCONST:       return "SCONST";
        case BCONST:       return "BCONST";
        case XCONST:       return "XCONST";
        case USCONST:      return "USCONST";
        case PARAM:        return "PARAM";
        case IDENT:        return "IDENT";
        case UIDENT:       return "UIDENT";
        case OP:           return "Op";
        case TYPECAST:     return "TYPECAST";
        case DOT_DOT:      return "DOT_DOT";
        case COLON_EQUALS: return "COLON_EQUALS";
        case EQUALS_GREATER: return "EQUALS_GREATER";
        case LESS_EQUALS:  return "LESS_EQUALS";
        case GREATER_EQUALS: return "GREATER_EQUALS";
        case NOT_EQUALS:   return "NOT_EQUALS";
        case RIGHT_ARROW:  return "RIGHT_ARROW";
        case FORMAT_LA:    return "FORMAT_LA";
        case NOT_LA:       return "NOT_LA";
        case NULLS_LA:     return "NULLS_LA";
        case WITH_LA:      return "WITH_LA";
        case WITHOUT_LA:   return "WITHOUT_LA";
        /* Single-character tokens (named in Lime) */
        case LPAREN:       return "(";
        case RPAREN:       return ")";
        case LBRACKET:     return "[";
        case RBRACKET:     return "]";
        case COMMA:        return ",";
        case SEMICOLON:    return ";";
        case COLON:        return ":";
        case DOT:          return ".";
        case PLUS:         return "+";
        case MINUS:        return "-";
        case STAR:         return "*";
        case SLASH:        return "/";
        case PERCENT:      return "%";
        case CARET:        return "^";
        case LT:           return "<";
        case GT:           return ">";
        case EQ:           return "=";
        case PIPE:         return "|";
        case HASH:         return "#";
        default: {
            const char *name = pg_keyword_name(token);
            if (name) return name;
            return "???";
        }
    }
}


/* ======================================================================
** Lookahead token resolution
**
** PostgreSQL's parser.c performs one-token lookahead to resolve certain
** grammar ambiguities.  Five keywords are replaced with _LA variants
** depending on what follows:
**
**   NOT     -> NOT_LA      if next is BETWEEN, IN, LIKE, ILIKE, SIMILAR
**   NULLS   -> NULLS_LA    if next is FIRST or LAST
**   WITH    -> WITH_LA     if next is TIME or ORDINALITY
**   WITHOUT -> WITHOUT_LA  if next is TIME
**   FORMAT  -> FORMAT_LA   if next is JSON
**
** This allows the grammar to give NOT_LA the same precedence as LIKE
** (rather than NOT's normal precedence), and similarly for the others.
** ====================================================================== */

int pg_scan_with_lookahead(PgScanner *s, PgTokenValue *val, int *loc) {
    int token;

    /* If we have a buffered lookahead token, return it */
    if (s->have_lookahead) {
        s->have_lookahead = false;
        *val = s->la_val;
        *loc = s->la_loc;
        return s->la_token;
    }

    /* Get the current token */
    token = pg_scan(s, val, loc);
    if (token <= 0) return token;

    /* Check if this token needs lookahead resolution */
    bool needs_lookahead = (token == NOT ||
                               token == NULLS_P ||
                               token == WITH ||
                               token == WITHOUT ||
                               token == FORMAT);

    if (!needs_lookahead) return token;

    /* Peek at the next token */
    int next_token = pg_scan(s, &s->la_val, &s->la_loc);
    if (next_token <= 0) {
        /* EOF or error -- no replacement; push back nothing */
        /* But we need to preserve the EOF/error state */
        if (next_token == 0) {
            /* EOF: just return the original token, no lookahead buffered */
            return token;
        }
        return next_token;  /* propagate error */
    }

    /* Buffer the lookahead token */
    s->la_token = next_token;
    s->have_lookahead = true;

    /* Now check for each replacement rule */
    if (token == NOT) {
        if (next_token == BETWEEN || next_token == IN_P ||
            next_token == LIKE || next_token == ILIKE ||
            next_token == SIMILAR) {
            return NOT_LA;
        }
    }
    else if (token == NULLS_P) {
        if (next_token == FIRST_P || next_token == LAST_P) {
            return NULLS_LA;
        }
    }
    else if (token == WITH) {
        if (next_token == TIME || next_token == ORDINALITY) {
            return WITH_LA;
        }
    }
    else if (token == WITHOUT) {
        if (next_token == TIME) {
            return WITHOUT_LA;
        }
    }
    else if (token == FORMAT) {
        if (next_token == JSON) {
            return FORMAT_LA;
        }
    }

    /* No replacement needed */
    return token;
}


/* ======================================================================
** Standalone test driver
**
** Compile with: cc -o pg_tokenize tokenize.c -DPG_TOKENIZE_MAIN
** Usage: echo "SELECT * FROM foo;" | ./pg_tokenize
** ====================================================================== */

#ifdef PG_TOKENIZE_MAIN

static char *read_stdin(size_t *out_len) {
    size_t alloc = 4096;
    size_t len = 0;
    char *buf = malloc(alloc);
    if (!buf) return NULL;
    size_t n;
    while ((n = fread(buf + len, 1, alloc - len, stdin)) > 0) {
        len += n;
        if (len >= alloc) {
            alloc *= 2;
            char *newbuf = realloc(buf, alloc);
            if (!newbuf) { free(buf); return NULL; }
            buf = newbuf;
        }
    }
    buf[len] = '\0';
    *out_len = len;
    return buf;
}

int main(void) {
    size_t len;
    char *input = read_stdin(&len);
    if (!input) {
        fprintf(stderr, "Failed to read input\n");
        return 1;
    }

    PgScanner *scanner = pg_scanner_create(input, len);
    if (!scanner) {
        fprintf(stderr, "Failed to create scanner\n");
        free(input);
        return 1;
    }

    PgTokenValue val;
    int loc;
    int token;

    while ((token = pg_scan_with_lookahead(scanner, &val, &loc)) != 0) {
        if (token == -1) {
            fprintf(stderr, "Error at byte %d: %s\n",
                    pg_scanner_error_location(scanner),
                    pg_scanner_error(scanner));
            break;
        }

        printf("%-20s loc=%-5d", pg_token_name(token), loc);

        if (token == ICONST || token == PARAM) {
            printf(" val=%d", val.ival);
        } else if (token == FCONST || token == SCONST ||
                   token == BCONST || token == XCONST ||
                   token == USCONST || token == IDENT ||
                   token == UIDENT || token == OP) {
            printf(" val=\"%s\"", val.str);
            free(val.str);
        } else {
            const char *kwn = pg_keyword_name(token);
            if (kwn)
                printf(" kw=\"%s\"", val.keyword);
        }

        printf("\n");
    }

    pg_scanner_destroy(scanner);
    free(input);
    return (token == -1) ? 1 : 0;
}

#endif /* PG_TOKENIZE_MAIN */