from dataclasses import dataclass

from lexer import Identifier, Punctuator, IntegerConstant

@dataclass
class Tree:
    operation: str
    children: list

class Parser:
    def __init__(self, tokens):
        self.pos = 0
        self.tokens = tokens

    def advance(self, offset=1):
        token = self.tokens[self.pos]
        self.pos += offset
        return token

    def check(self, t, offset=0):
        if self.pos + offset < len(self.tokens):
            token = self.tokens[self.pos + offset]
            return type(token) is t
        else:
            return False

    def check_punctuator(self, p, offset=0):
        if self.pos + offset < len(self.tokens):
            token = self.tokens[self.pos + offset]
            return type(token) is Punctuator and token.token == p
        else:
            return False

    def check_identifier(self, p, offset=0):
        if self.pos + offset < len(self.tokens):
            token = self.tokens[self.pos + offset]
            return type(token) is Identifier and token.token == p
        else:
            return False

    def match(self, t):
        ret = self.check(t)
        if ret:
            self.advance()
        return ret

    def match_punctuator(self, p):
        ret = self.check_punctuator(p)
        if ret:
            self.advance()
        return ret

    def match_identifier(self, p):
        ret = self.check_identifier(p)
        if ret:
            self.advance()
        return ret

    def previous(self):
        return self.tokens[self.pos - 1]

    def primary_expression(self):
        aa = ' '.join(t.token for t in self.tokens[self.pos:])
        if self.match(Identifier):
            return self.previous()
        elif self.match(IntegerConstant):
            return self.previous()
        elif self.match_punctuator("("):
            expr = self.expression()
            assert self.match_punctuator(")"), self.advance()
            return Tree(
                operation="grouping",
                children=[expr]
            )
        else:
            assert False, self.advance()

    def argument_expression_list(self):
        expr = self.assignment_expression()
        while True:
            if self.match_punctuator(","):
                right = self.assignment_expression()
                expr = Tree(
                    operation="argument_list",
                    children=[expr, right]
                )
            else:
                break
        return expr

    def for_expression(self):
        left = self.primary_expression()
        assert self.match_identifier("FOR"), self.advance()
        right = self.primary_expression()
        return Tree(
            operation="for_expression",
            children=[left, right]
        )

    def postfix_expression(self):
        expr = self.primary_expression()
        while True:
            if self.match_punctuator("["):
                right = self.expression()
                assert self.match_punctuator("]"), self.advance()
                expr = Tree(
                    operation="subscript",
                    children=[expr, right]
                )
            elif self.match_punctuator("("):
                if self.match_punctuator(")"):
                    right = []
                else:
                    right = [self.argument_expression_list()]
                    assert self.match_punctuator(")"), self.advance()
                expr = Tree(
                    operation="function_call",
                    children=[expr, *right]
                )
            elif self.match_punctuator("<"):
                backtrack = self.pos - 1
                try:
                    right = self.for_expression()
                except AssertionError:
                    self.pos = backtrack
                    break
                assert self.match_punctuator(">"), self.advance()
                expr = Tree(
                    operation="bit_extraction",
                    children=[expr, right]
                )
            elif self.match_punctuator("."):
                assert self.match(Identifier), self.advance()
                right = self.previous()
                expr = Tree(
                    operation="member",
                    children=[expr, right]
                )
            else:
                break
        return expr

    def unary_expression(self):
        if self.match_identifier("NOT"):
            return Tree(
                operation="unary_not",
                children = [
                    self.unary_expression()
                ]
            )
        elif self.match_identifier("INT"):
            return Tree(
                operation="unary_int",
                children = [
                    self.unary_expression()
                ]
            )
        elif self.match_punctuator("~"):
            return Tree(
                operation="unary_complement",
                children = [
                    self.unary_expression()
                ]
            )
        elif self.match_punctuator("-"):
            return Tree(
                operation="unary_negation",
                children = [
                    self.unary_expression()
                ]
            )
        # elif self.match_punctuator("|"):
        #     expr = self.unary_expression()
        #     assert self.match_punctuator("|"), self.advance()
        #     return Tree(
        #         operation="unary_absolute_value",
        #         children = [
        #             expr
        #         ]
        #     )
        else:
            return self.postfix_expression()

    def multiplicative_expression(self):
        expr = self.unary_expression()
        while self.match_punctuator("×") or self.match_punctuator("/"):
            operation = {
                "×": "multiplication",
                "/": "division",
            }[self.previous().token]
            expr = Tree(
                operation=operation,
                children=[
                    expr,
                    self.unary_expression()
                ]
            )
        return expr

    def additive_expression(self):
        expr = self.multiplicative_expression()
        while self.match_punctuator("+") or self.match_punctuator("-"):
            operation = {
                "+": "addition",
                "-": "subtraction",
            }[self.previous().token]
            expr = Tree(
                operation=operation,
                children=[
                    expr,
                    self.multiplicative_expression()
                ]
            )
        return expr

    def shift_expression(self):
        expr = self.additive_expression()
        while self.match_punctuator("<<") or self.match_punctuator(">>"):
            operation = {
                "<<": "left_shift",
                ">>": "right_shift",
            }[self.previous().token]
            expr = Tree(
                operation=operation,
                children=[
                    expr,
                    self.additive_expression()
                ]
            )
        return expr

    def relational_expression(self):
        expr = self.shift_expression()
        while self.match_punctuator("<") or self.match_punctuator(">") \
              or self.match_punctuator("≤") or self.match_punctuator("≥"):
            operation = {
                "<": "less_than",
                ">": "greater_than",
                "≤": "less_than_equal",
                "≥": "greater_than_equal",
            }[self.previous().token]
            expr = Tree(
                operation=operation,
                children=[
                    expr,
                    self.shift_expression()
                ]
            )
        return expr
    
    def equality_expression(self):
        expr = self.relational_expression()
        while self.match_punctuator("=") or self.match_punctuator("≠"):
            operation = {
                "≠": "equality_not_equal",
                "=": "equality_equal",
            }[self.previous().token]
            expr = Tree(
                operation=operation,
                children=[
                    expr,
                    self.relational_expression()
                ]
            )
        return expr

    def bitwise_and_expression(self):
        expr = self.equality_expression()
        while self.match_punctuator("∧"):
            expr = Tree(
                operation="bitwise_and",
                children=[
                    expr,
                    self.equality_expression()
                ]
            )
        return expr

    def bitwise_xor_expression(self):
        expr = self.bitwise_and_expression()
        while self.match_punctuator("⊕"):
            expr = Tree(
                operation="bitwise_xor",
                children=[
                    expr,
                    self.bitwise_and_expression()
                ]
            )
        return expr

    def bitwise_or_expression(self):
        expr = self.bitwise_xor_expression()
        while self.match_punctuator("∨"):
            expr = Tree(
                operation="bitwise_or",
                children=[
                    expr,
                    self.bitwise_xor_expression()
                ]
            )
        return expr

    def logical_and_expression(self):
        expr = self.bitwise_or_expression()
        while self.match_identifier('AND'):
            expr = Tree(
                operation="logical_and",
                children=[
                    expr,
                    self.bitwise_or_expression()
                ]
            )
        return expr

    def logical_xor_expression(self):
        # expr = self.logical_and_expression()
        # while self.match_identifier('XOR'):
        #     expr = Tree(
        #         operation="logical_xor",
        #         children=[
        #             expr,
        #             self.logical_and_expression()
        #         ]
        #     )
        # return expr
        return self.logical_and_expression()

    def logical_or_expression(self):
        expr = self.logical_xor_expression()
        while self.match_identifier('OR'):
            expr = Tree(
                operation="logical_or",
                children=[
                    expr,
                    self.logical_xor_expression()
                ]
            )
        return expr

    def assignment_list(self):
        expr = self.unary_expression()
        while True:
            if self.match_punctuator(","):
                right = self.unary_expression()
                expr = Tree(
                    operation="assignment_list",
                    children=[expr, right]
                )
            else:
                break
        return expr

    def assignment_expression(self):
        backtrack = self.pos
        try:
            left = self.assignment_list()
        except AssertionError:
            self.pos = backtrack
            return self.logical_or_expression()

        if self.match_punctuator("←"):
            right = self.assignment_expression()
            return Tree(
                operation="assignment",
                children=[left, right],
            )
        else:
            self.pos = backtrack
            return self.logical_or_expression()

    def expression(self):
        return self.assignment_expression()

    def statement(self):
        return self.unlabeled_statement()

    def unlabeled_statement(self):
        backtrack = self.pos
        try:
            return self.expression_statement()
        except AssertionError:
            self.pos = backtrack
            return self.primary_block()

    def expression_statement(self):
        expr = self.expression()
        assert self.match_punctuator(";"), self.advance()
        return Tree(
            operation="expression_statement",
            children=[expr]
        )

    def primary_block(self):
        if self.check_punctuator("{"):
            return self.compound_statement()
        elif self.check_identifier("IF"):
            return self.selection_statement()
        elif self.check_identifier("THROW"):
            return self.throw_statement()
        else:
            assert False, self.advance()

    def secondary_block(self):
        return self.statement()

    def compound_statement(self):
        assert self.match_punctuator("{"), self.advance()
        if self.match_punctuator("}"):
            return Tree(
                operation="compound_statement",
                children=[],
            )
        else:
            expr = self.block_item_list()
            assert self.match_punctuator("}"), self.advance()
            return Tree(
                operation="compound_statement",
                children=[expr]
            )

    def block_item_list(self):
        expr = self.block_item()
        while True:
            backtrack = self.pos
            try:
                right = self.block_item()
                expr = Tree(
                    operation="block_item_list",
                    children=[expr, right]
                )
            except AssertionError:
                self.pos = backtrack
                break
        return expr

    def block_item(self):
        return self.unlabeled_statement()

    def selection_statement(self):
        assert self.match_identifier("IF"), self.advance()
        assert self.match_punctuator("("), self.advance()
        cond_expr = self.expression()
        assert self.match_punctuator(")"), self.advance()
        true_block = self.secondary_block()
        if self.match_identifier("ELSE"):
            else_block = self.secondary_block()
            return Tree(
                operation="if_else",
                children=[cond_expr, true_block, else_block]
            )
        else:
            return Tree(
                operation="if",
                children=[cond_expr, true_block]
            )

    def throw_statement(self):
        assert self.match_identifier("THROW"), self.advance()
        block = self.primary_expression()
        if self.match_punctuator(","):
            arg = self.primary_expression()
            assert self.match_punctuator(";"), self.advance()
            return Tree(
                operation="throw_arg",
                children=[block, arg]
            )
        else:
            assert self.match_punctuator(";"), self.advance()
            return Tree(
                operation="throw",
                children=[block]
            )