from lark import Lark
from lark.visitors import Visitor
from lark.lexer import Token

class ACPBench_Visitor(Visitor):
    def __init__(self) -> None:
        super().__init__()
        self.action_lists = None
        self.action_names = None
        self.progression_lists = None
        self.prog_lists = None
        self.indexes = None

    def action_list(self, tree):
        self.action_lists = []

    def prog_list(self, tree):
        if self.prog_lists is not None:
            self.progression_lists.append(self.prog_lists)
        self.prog_lists = []

    def progression_list(self, tree):
        self.progression_lists = []

    def action_none(self, tree):
        self.action_names = 'None'

    def action_name(self, tree):
        act_name = '(' + ''.join(tree.children[1:-1]) + ')'
        self.action_names = act_name
        if self.action_lists is not None:
            self.action_lists.append(act_name)
        if self.prog_lists is not None:
            self.prog_lists.append(act_name)

    def index(self, tree):
        self.indexes = ''.join(tree.children)
        if not self.indexes.isnumeric():
            self.indexes = None


class ACPGrammarParser(object):
    def __init__(self, grammarfile, task) -> None:
        self.task = task
        with open(grammarfile) as f:
            grammar = f.read()
            self.acp_parser = Lark(grammar, start=task, parser='lalr')

    def parse(self, input, debug=False):

        def ignore_errors(e):
            if hasattr(e, 'token') and e.token.type == '$END':
                for x in e.expected:
                    if x != 'WS':
                        e.interactive_parser.feed_token(Token(x,  self.acp_parser.get_terminal(x).pattern.value))
            
            return True
        

        input = input.replace('\n', '')
        input = input.strip()
        try:
            tree = self.acp_parser.parse(input, on_error=ignore_errors)
            
            if debug:
                print(tree)
            visitor = ACPBench_Visitor()
            visitor.visit_topdown(tree)
            if self.task == 'action_list':
                return visitor.action_lists
            elif self.task == 'act':
                return visitor.action_names
            elif self.task == 'action_name':
                return visitor.action_names
            elif self.task == 'index':
                return visitor.indexes
            elif self.task == 'progression_list':
                if visitor.prog_lists not in visitor.progression_lists:
                    visitor.progression_lists.append(visitor.prog_lists)
                return visitor.progression_lists
        except Exception as e:
            if debug:
                print('exception')
                print(e)
            return None


 


  
