#!/usr/bin/env python3
"""
Generate the keyword lookup table for the PostgreSQL Lime parser.

Reads:
  - pg_grammar.h   : Lime-generated token ID #defines
  - gram.lime      : keyword category productions (unreserved_keyword,
                     col_name_keyword, type_func_name_keyword,
                     reserved_keyword, bare_label_keyword)

Writes:
  - pg_kwlist.h    : Complete keyword table with token IDs, categories,
                     and bare-label flags, sorted alphabetically.
"""

import re
import sys
import os

def parse_token_ids(header_path):
    """Parse #define TOKEN_NAME value from pg_grammar.h."""
    token_ids = {}
    with open(header_path) as f:
        for line in f:
            m = re.match(r'^#define\s+(\w+)\s+(\d+)', line)
            if m:
                token_ids[m.group(1)] = int(m.group(2))
    return token_ids

def parse_keyword_category(gram_path, category_name):
    """Extract keyword tokens from a category production in gram.lime.

    e.g., unreserved_keyword(A) ::= ABORT_P.
    """
    keywords = set()
    pattern = re.compile(
        r'^' + re.escape(category_name) + r'\([A-Z]\)\s*::=\s*(\w+)\.'
    )
    with open(gram_path) as f:
        for line in f:
            m = pattern.match(line)
            if m:
                keywords.add(m.group(1))
    return keywords

def token_to_keyword_string(token_name):
    """Convert token name to lowercase keyword string.

    ABORT_P -> "abort"
    ABSOLUTE_P -> "absolute"
    CHAR_P -> "char"
    IN_P -> "in"
    """
    s = token_name
    # Strip trailing _P suffix (PostgreSQL convention for C reserved words)
    if s.endswith('_P'):
        s = s[:-2]
    return s.lower()

def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    header_path = os.path.join(script_dir, 'pg_grammar.h')
    gram_path = os.path.join(script_dir, 'gram.lime')
    output_path = os.path.join(script_dir, 'pg_kwlist.h')

    if not os.path.exists(header_path):
        print(f"Error: {header_path} not found. Run Lime first.", file=sys.stderr)
        sys.exit(1)
    if not os.path.exists(gram_path):
        print(f"Error: {gram_path} not found.", file=sys.stderr)
        sys.exit(1)

    # Parse token IDs
    token_ids = parse_token_ids(header_path)
    print(f"Loaded {len(token_ids)} token IDs from pg_grammar.h")

    # Parse keyword categories
    unreserved = parse_keyword_category(gram_path, 'unreserved_keyword')
    col_name = parse_keyword_category(gram_path, 'col_name_keyword')
    type_func = parse_keyword_category(gram_path, 'type_func_name_keyword')
    reserved = parse_keyword_category(gram_path, 'reserved_keyword')
    bare_label = parse_keyword_category(gram_path, 'bare_label_keyword')

    print(f"  unreserved_keyword:    {len(unreserved)}")
    print(f"  col_name_keyword:      {len(col_name)}")
    print(f"  type_func_name_keyword:{len(type_func)}")
    print(f"  reserved_keyword:      {len(reserved)}")
    print(f"  bare_label_keyword:    {len(bare_label)}")

    # All keywords = union of the four categories
    all_keywords = unreserved | col_name | type_func | reserved

    # Verify all keywords have token IDs
    missing = []
    for kw in sorted(all_keywords):
        if kw not in token_ids:
            missing.append(kw)
    if missing:
        print(f"WARNING: {len(missing)} keywords missing from pg_grammar.h:")
        for m in missing:
            print(f"  {m}")

    # Build sorted keyword entries
    entries = []
    for token_name in sorted(all_keywords, key=lambda t: token_to_keyword_string(t)):
        kw_string = token_to_keyword_string(token_name)
        token_id = token_ids.get(token_name, 0)

        if token_name in unreserved:
            cat = 'PG_KW_UNRESERVED'
        elif token_name in col_name:
            cat = 'PG_KW_COL_NAME'
        elif token_name in type_func:
            cat = 'PG_KW_TYPE_FUNC_NAME'
        elif token_name in reserved:
            cat = 'PG_KW_RESERVED'
        else:
            cat = 'PG_KW_UNRESERVED'  # fallback

        is_bare = 'true' if token_name in bare_label else 'false'

        entries.append((kw_string, token_name, token_id, cat, is_bare))

    # Generate the header file
    with open(output_path, 'w') as f:
        f.write("""\
/*
 * pg_kwlist.h
 *    PostgreSQL keyword lookup table for Lime parser.
 *
 * AUTO-GENERATED by gen_kwlist.py -- do not edit manually.
 *
 * Each entry maps a lowercase keyword string to its Lime token ID,
 * reservation category, and bare-label eligibility.
 */

#ifndef PG_KWLIST_H
#define PG_KWLIST_H

#include <stdbool.h>
#include <stddef.h>
#include <string.h>
#include "pg_grammar.h"  /* Token ID #defines */

typedef enum {
    PG_KW_UNRESERVED,     /* Can be used as any identifier */
    PG_KW_COL_NAME,       /* Can be column/table name, not type/func name */
    PG_KW_TYPE_FUNC_NAME, /* Can be type or function name */
    PG_KW_RESERVED        /* Only usable as ColLabel (after AS) */
} PgKwCategory;

typedef struct {
    const char *keyword;     /* Lowercase keyword string */
    int         token_id;    /* Lime token ID from pg_grammar.h */
    PgKwCategory category;   /* Reservation category */
    bool        is_bare_label; /* Can be used as column label without AS */
} PgKwEntry;

""")

        f.write(f"#define PG_NUM_KEYWORDS {len(entries)}\n\n")

        f.write("static const PgKwEntry pg_keywords[] = {\n")
        for kw_string, token_name, token_id, cat, is_bare in entries:
            # Use the #define name for token_id for clarity and correctness
            f.write(f'    {{"{kw_string}",'
                    f' {token_name:30s}, {cat:25s}, {is_bare:5s}}},\n')
        f.write("};\n\n")

        f.write("""\
/*
 * Binary search for a keyword. Input must be lowercase.
 * Returns a pointer to the PgKwEntry, or NULL if not found.
 */
static inline const PgKwEntry *
pg_kw_lookup(const char *str, size_t len)
{
    int lo = 0;
    int hi = PG_NUM_KEYWORDS - 1;

    while (lo <= hi) {
        int mid = (lo + hi) / 2;
        const PgKwEntry *entry = &pg_keywords[mid];
        int cmp = strncmp(str, entry->keyword, len);
        if (cmp == 0) {
            /* Check for exact length match */
            if (entry->keyword[len] == '\\0')
                return entry;
            /* str is a prefix of keyword -- str is "less" */
            cmp = -1;
        }
        if (cmp < 0)
            hi = mid - 1;
        else
            lo = mid + 1;
    }
    return NULL;
}

#endif /* PG_KWLIST_H */
""")

    print(f"\nGenerated {output_path} with {len(entries)} keywords.")

if __name__ == '__main__':
    main()