
from string import ascii_letters, digits
from rply import LexerGenerator

VALID_CHARACTERS = ascii_letters+"_"+digits

TOKENS = [
    ("ROLL_DIE",            r"d"),
    ("ROLL_KEEP_LOWEST",    r"kl"),
    ("ROLL_KEEP_HIGHEST",   r"kh?"),
    ("ROLL_REROLL_ONCE",    r"ro"),
    ("ROLL_REROLL",         r"r"),
    ("ROLL_MIN",            r"min"),
    ("ROLL_MAX",            r"max"),
    ("ROLL_EXPLODE",        r"!"),

    ("SYMBOL_LE",           r"\<\="),
    ("SYMBOL_GE",           r"\>\="),
    ("SYMBOL_LT",           r"\<"),
    ("SYMBOL_GT",           r"\>"),
    ("SYMBOL_EQUALS",       r"\="),
    ("SYMBOL_ARROW",        r"\-\>"),
    ("SYMBOL_BACKSLASH",    r"\\"),
    ("SYMBOL_LPARENS",      r"\("),
    ("SYMBOL_RPARENS",      r"\)"),
    ("SYMBOL_PLUS",         r"\+"),
    ("SYMBOL_MINUS",        r"\-"),
    ("SYMBOL_TIMES",        r"\*"),
    ("SYMBOL_DIVIDE",       r"\/"),

    ("KEYWORD_LET",         r"let"),
    ("KEYWORD_IN",          r"in"),
    ("KEYWORD_IF",          r"if"),
    ("KEYWORD_THEN",        r"then"),
    ("KEYWORD_ELSE",        r"else"),

    ("INT",                 r"\d+"),
    ("ID",                  f"[{ascii_letters}][{VALID_CHARACTERS}]*")
]

class Lexer():
    def __init__(self):
        self.lexer = LexerGenerator()

    def _add_tokens(self):
        for token in TOKENS:
            self.lexer.add(*token)
        self.lexer.ignore(r"\s+")

    def get_lexer(self):
        self._add_tokens()
        return self.lexer.build()
