from rply import ParserGenerator

from lexer import TOKENS
from ast_nodes import Exp, ExpInt, ExpBinop, ExpLet, ExpVar, ExpMin, ExpMax, ExpRoll, Roll, RollKeepHighest, RollKeepLowest, RollMin, RollMax, RollExplode, ComparePoint, ExpIf, ExpTest, ExpApply, ExpLambda, ExpNeg, RollReroll


class Parser():
    def __init__(self):
        self.pg = ParserGenerator(
            [i[0] for i in TOKENS],
            precedence=[
                ('left', ["SYMBOL_BACKSLASH","SYMBOL_ARROW"]),
                ('left', ["KEYWORD_LET", "KEYWORD_IN", "KEYWORD_IF", "KEYWORD_THEN", "KEYWORD_ELSE"]),
                ('left', ["SYMBOL_EQUALS", "SYMBOL_LT", "SYMBOL_LE", "SYMBOL_GT", "SYMBOL_GE"]),
                ('left', ["SYMBOL_PLUS", "SYMBOL_MINUS"]),
                ('left', ["SYMBOL_TIMES", "SYMBOL_DIVIDE"]),
                ('left', ["ROLL_DIE"]),
                ('right', ["ROLL_KEEP_HIGHEST","ROLL_KEEP_LOWEST","ROLL_MIN","ROLL_MAX","ROLL_EXPLODE","ROLL_REROLL","ROLL_REROLL_ONCE"])
            ]
        )
        self._get_parser()

    def parse(self, token_input):
        return self.parser.parse(token_input)

    def _get_parser(self) -> Exp:
        # Expressions
        @self.pg.production('exp : exp SYMBOL_PLUS exp')
        @self.pg.production('exp : exp SYMBOL_MINUS exp')
        @self.pg.production('exp : exp SYMBOL_TIMES exp')
        @self.pg.production('exp : exp SYMBOL_DIVIDE exp')
        def exp_mul_div(tokens):
            return ExpBinop(tokens[1].value, tokens[0], tokens[2])

        @self.pg.production('exp : atom')
        def exp_atom(tokens):
            return tokens[0]

        @self.pg.production("exp : SYMBOL_MINUS atom")
        def exp_neg(tokens):
            return ExpNeg(tokens[1])

        @self.pg.production('exp : KEYWORD_LET ID SYMBOL_EQUALS exp KEYWORD_IN exp')
        def exp_let(tokens):
            return ExpLet(tokens[1].value, tokens[3], tokens[5])

        @self.pg.production('exp : roll')
        def exp_roll(tokens):
            return ExpRoll(tokens[0])

        @self.pg.production("exp : atom ROLL_MIN atom")
        def exp_min(tokens):
            return ExpMin(tokens[0], tokens[2])

        @self.pg.production("exp : atom ROLL_MAX atom")
        def exp_max(tokens):
            return ExpMax(tokens[0], tokens[2])

        @self.pg.production("exp : atom comp")
        def exp_test(tokens):
            return ExpTest(tokens[0], tokens[1])

        @self.pg.production("exp : KEYWORD_IF exp KEYWORD_THEN exp KEYWORD_ELSE exp")
        def exp_if(tokens):
            return ExpIf(tokens[1],tokens[3],tokens[5])

        @self.pg.production("exp : SYMBOL_BACKSLASH ID SYMBOL_ARROW exp")
        def exp_lamda(tokens):
            return ExpLambda(tokens[1].value, tokens[3])

        @self.pg.production("exp : atom atom")
        def apply(tokens):
            return ExpApply(tokens[0],tokens[1])

        # Rolls
        @self.pg.production('roll : roll ROLL_KEEP_HIGHEST atom')
        def roll_keep_highest(tokens):
            return RollKeepHighest(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_KEEP_LOWEST atom')
        def roll_keep_lowest(tokens):
            return RollKeepLowest(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_MIN atom')
        def roll_min(tokens):
            return RollMin(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_MAX atom')
        def roll_max(tokens):
            return RollMax(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_EXPLODE ')
        def roll_explode(tokens):
            return RollExplode(tokens[0])

        @self.pg.production('roll : roll ROLL_EXPLODE comp')
        def roll_explode_comp(tokens):
            return RollExplode(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_REROLL ')
        def roll_reroll(tokens):
            return RollReroll(tokens[0])

        @self.pg.production('roll : roll ROLL_REROLL comp')
        def roll_reroll_comp(tokens):
            return RollReroll(tokens[0], tokens[2])

        @self.pg.production('roll : roll ROLL_REROLL_ONCE')
        def roll_reroll_once(tokens):
            return RollReroll(tokens[0], None, True)

        @self.pg.production('roll : roll ROLL_REROLL_ONCE comp')
        def roll_reroll_once_comp(tokens):
            return RollReroll(tokens[0], tokens[2], True)

        @self.pg.production('roll : atom ROLL_DIE atom')
        def roll(tokens):
            return Roll(tokens[0], tokens[2])

        @self.pg.production('roll : ROLL_DIE atom')
        def roll_no_amount(tokens):
            return Roll(ExpInt(1), tokens[1])

        # Compare Points

        @self.pg.production("comp : SYMBOL_EQUALS atom")
        @self.pg.production("comp : SYMBOL_LT atom")
        @self.pg.production("comp : SYMBOL_LE atom")
        @self.pg.production("comp : SYMBOL_GT atom")
        @self.pg.production("comp : SYMBOL_GE atom")
        def comp_point(tokens):
            return ComparePoint(tokens[0].value,tokens[1])

        # Atoms
        @self.pg.production("atom : INT")
        def atom_int(tokens):
            return ExpInt(int(tokens[0].value))

        @self.pg.production("atom : SYMBOL_LPARENS exp SYMBOL_RPARENS")
        def atom_exp(tokens):
            return tokens[1]

        @self.pg.production('atom : ID')
        def atom_var(tokens):
            return ExpVar(tokens[0].value)

        ## Error Handling ##
        @self.pg.error
        def error_handle(token):
            raise Exception(f"Unexpected token '{token.value}' ({token.name}) at line {token.source_pos.lineno}, column {token.source_pos.colno}.")

        ## Finish ##
        self.parser = self.pg.build()