from fol.node import Node


class Parser:
    def __init__(self, lexer):
        self.lexer = lexer

    def consume(self, name):
        tok = self.lexer.peek()

        if tok and tok.is_a(name):
            self.lexer.next()
            return tok

        return None

    def ensure(self, name):
        tok = self.consume(name)

        if tok is None:
            raise Exception(f'invalid {name}')

        if tok and not tok.is_a(name):
            raise Exception(f'expected {name}, got {tok.name}')

        return tok

    def parse(self):
        return self.parse_f()

    def parse_f(self):
        return self.parse_exists()

    def parse_exists(self):
        if self.consume('EXISTS'):
            tok_var = self.ensure('VAR')
            rhs = self.parse_forall()
            node = Node('EXISTS')
            node.add_child(Node('VAR', tok_var.value))
            node.add_child(rhs)
            return node

        return self.parse_forall()

    def parse_forall(self):
        if self.consume('FORALL'):
            tok_var = self.ensure('VAR')
            rhs = self.parse_imp()
            node = Node('FORALL')
            node.add_child(Node('VAR', tok_var.value))
            node.add_child(rhs)
            return node

        return self.parse_imp()

    def parse_imp(self):
        lhs = self.parse_or()

        if self.consume('IMP'):
            rhs = self.parse_or()
            node = Node('IMP')
            node.add_child(lhs)
            node.add_child(rhs)
            lhs = node

        return lhs

    def parse_or(self):
        lhs = self.parse_and()

        while self.consume('OR'):
            rhs = self.parse_and()
            node = Node('OR')
            node.add_child(lhs)
            node.add_child(rhs)
            lhs = node

        return lhs

    def parse_and(self):
        lhs = self.parse_not()

        while self.consume('AND'):
            rhs = self.parse_not()
            node = Node('AND')
            node.add_child(lhs)
            node.add_child(rhs)
            lhs = node

        return lhs

    def parse_not(self):
        if self.consume('NOT'):
            node = Node('NOT')
            node.add_child(self.parse_group())
            return node

        return self.parse_group()

    def parse_group(self):
        if self.consume('OPAR'):
            node = self.parse_f()
            self.ensure('CPAR')
            return node
        return self.parse_pred()

    def parse_pred(self):
        tok = self.ensure('PRED')
        self.ensure('OPAR')

        node = Node('PRED', tok.value)
        node.add_child(self.parse_term())

        while self.consume('COMMA'):
            node.add_child(self.parse_term())

        self.ensure('CPAR')
        return node

    def parse_term(self):
        maybe_var = self.lexer.peek(0)
        maybe_opar = self.lexer.peek(1)

        if maybe_var and maybe_opar and maybe_var.is_a('VAR') \
           and maybe_opar.is_a('OPAR'):
            return self.parse_function()

        tok = self.lexer.next()
        if tok is None:
            raise Exception('unexpected end')

        if tok.name in ['VAR', 'CONST']:
            return Node(tok.name, tok.value)

        raise Exception(f'unknown term {tok.value}')

    def parse_function(self):
        tok = self.ensure('VAR')
        self.ensure('OPAR')

        node = Node('FUN', tok.value)
        node.add_child(self.parse_term())

        while self.consume('COMMA'):
            node.add_child(self.parse_term())

        self.ensure('CPAR')
        return node