"""Creates GeneratedParser's from the given grammar

Resources:
    https://youtu.be/QppWTvh7_sI
    https://en.wikipedia.org/wiki/Parsing_expression_grammar
    https://medium.com/@gvanrossum_83706/peg-parsing-series-de5d41b2ed60

"""
from typing import *

from .lexer import TokenStream
from .parser import *
from pprint import pprint
from contextlib import contextmanager


HEADER = """# Code @generated by parsergen; do not edit!
from parsergen.parser_utils import GeneratedParser, TokenStream, Node, Filler
from parsergen.parser_utils import memoize, memoize_left_rec
from functools import reduce
"""

class Generator:
    def __init__(self) -> None:
        self.result = ""
        self._indent = 0
        self.counter = 0
        self.queue = []
        self.config = {
            "class_name": "CustomParser",
            "inherits_from": "GeneratedParser",
            "header": "",
        }
    
    def push(self, line: str):
        lines = line.split("\n")
        for l in lines:
            self.result += (self._indent * "    ") + l + "\n"
    
    @contextmanager
    def indent(self):
        saved = self._indent
        try:
            self._indent += 1
            yield
        finally:
            self._indent = saved
    
    def statement_mapping_graph(self, rules: Dict[str, List[Statement]]):
        mapping = {}
        for name, rs in rules.items():
            mapping[name] = set()
            for r in rs:
                mapping[name] |= self.map_node(r)
        return mapping

    def map_node(self, node) -> Set[str]:
        result = set()
        if isinstance(node, Statement):
            for part in node.grammar:
                result |= self.map_node(part)
        if isinstance(node, ExprList):
            for part in node.exprs:
                result |= self.map_node(part)
        if isinstance(node, (Quantifier, NamedItem)):
            result |= self.map_node(node.expr)
        if isinstance(node, StatementPointer):
            result |= {node.target}
        return result
    
    def is_left_recursive(self, grammar, name: str):
        data = self.statement_mapping_graph(grammar)
        locations = self.located_in(data, name)
        while True:
            current = set()
            prev = locations.copy()
            for loc in locations:
                current |= self.located_in(data, loc)
            locations |= current
            for loc in data[name]:
                if loc in locations:
                    return True
            if locations == prev:
                break
            
        return False

    def located_in(self, data: Dict[str, Set[str]], name: str) -> Set[str]:
        return {k for k in data if name in data[k]}
    
    def process_sections(self, parser_definition: ParserDefinition) -> Dict[str, List[Statement]]:
        rules_list = []
        for section in parser_definition.sections:
            if isinstance(section, Statement):
                rules_list.append(section)
            elif isinstance(section, ConfigurationCall):
                self.handle_ConfigurationCall(section)
        
        return post_process(rules_list)

    def generate(self, grammar: str, display=True):
        gl = GrammarLexer()
        parser_definition = parse_all(gl.lex_string(grammar))
        rules = self.process_sections(parser_definition)
        #pprint(rules)
        if display:
            print("Formatted grammar:")
            print(format_grammar(rules)) # TODO: change to print entire parser_definition 
        return self.generate_parser_class(rules)
    
    def handle_ConfigurationCall(self, call: ConfigurationCall):
        self.config[call.name] = call.value

    def generate_parser_class(self, grammar: Dict[str, List[Statement]]):
        
        self.push(f"class {self.config['class_name']}({self.config['inherits_from']}):")
        with self.indent():
            for name, stmts in grammar.items():
                if self.is_left_recursive(grammar, name):
                    self.push(f"@memoize_left_rec")
                else:
                    self.push(f"@memoize")
                self.push(f"def {name}(self):")
                queue = []
                with self.indent():
                    self.push(f"pos = self.mark()")
                    for stmt in stmts:
                        self.gen_statement(stmt, queue)
                    self.push(f"return None\n")
                self.resolve_queue(queue)
        return HEADER + self.config["header"] + self.result
    
    def gen_statement(self, stmt: Statement, queue):
        self.push(f'"""\n{GrammarPrinter(None).process(stmt)}\n"""')
        self.push("parts = []")
        named_items = []
        self.push("for _ in range(1):")
        with self.indent():
            for c, item in enumerate(stmt.grammar):
                if isinstance(item, NamedItem):
                    named_items.append((c, item))
                    item = item.expr
                self.gen(item, queue)
                if not isinstance(item, Predicate):
                    # don't generate extra clause for Predicates
                    self.push("if not self.match(part):")
                    with self.indent():
                        self.push("self.fail()")
                        self.push("break")
                    self.push("parts.append(part)")     
            self.push("# match:")
            for c, named_item in named_items:
                self.push(f"{named_item.name} = parts[{c}]")
            if stmt.action:
                self.push(f"return {stmt.action}")
            else:
                self.push(f"return Node({stmt.name!r}, parts)")  
        
        self.push(f"self.goto(pos)\n")
    
    def gen(self, item: Expr, queue):
        return getattr(self, f"gen_{type(item).__name__}")(item, queue)
    
    #def construct(self, f: str, args: List[str]):
    #    return f"{f}({', '.join(repr(arg) for arg in args)})"
    
    def gen_TokenPointer(self, item: TokenPointer, queue):
        self.push(f"part = self.expect({item.target!r})")
    
    def gen_StatementPointer(self, item: StatementPointer, queue):
        self.push(f"part = self.{item.target}()")
    
    def gen_ConstantString(self, item: ConstantString, queue):
        self.push(f"part = self.expect_constant({item.value!r})")
    
    def gen_AndPredicate(self, item: AndPredicate, queue):
        self.push("predicate_pos = self.mark()")
        self.gen(item.expr, queue)
        self.push("self.goto(predicate_pos)")
        self.push("if not self.match(part):")
        with self.indent():
            self.push("self.fail()")
            self.push("break")
    
    def gen_NotPredicate(self, item: NotPredicate, queue):
        self.push("predicate_pos = self.mark()")
        self.gen(item.expr, queue)
        self.push("self.goto(predicate_pos)")
        self.push("if self.match(part):")
        with self.indent():
            self.push("self.fail()")
            self.push("break")

    
    def gen_ZeroOrMore(self, item: ZeroOrMore, queue):
        queue.append((item, self.counter))
        self.push(f"part = self._loop_{self.counter}()")
        self.counter += 1
    
    def gen_OneOrMore(self, item: OneOrMore, queue):
        queue.append((item, self.counter))
        self.push(f"part = self._loop_{self.counter}()")
        self.counter += 1
    
    def gen_ZeroOrOne(self, item: ZeroOrOne, queue):
        queue.append((item, self.counter))
        self.push(f"part = self._maybe_{self.counter}()")
        self.counter += 1
    
    def gen_OrOp(self, item: OrOp, queue):
        queue.append((item, self.counter))
        self.push(f"part = self._or_{self.counter}()")
        self.counter += 1

    def gen_ExprList(self, item: ExprList, queue):
        queue.append((item, self.counter))
        self.push(f"part = self._expr_list_{self.counter}()")
        self.counter += 1
    
    def resolve(self, item: Expr, c, queue):
        return getattr(self, f"resolve_{type(item).__name__}")(item, c, queue)
    
    def resolve_ZeroOrMore(self, item: ZeroOrMore, c, queue):
        self.push(f"def _loop_{c}(self):")
        with self.indent():
            self.push(f'"""\n{GrammarPrinter(None).process(item)}\n"""')
            self.push("children = []")
            self.push("while True:")
            with self.indent():
                self.push("pos = self.mark()")
                self.gen(item.expr, queue)
                self.push("if self.match(part): children.append(part)")
                self.push("else:")
                with self.indent():
                    self.push("self.goto(pos)")
                    self.push("break")
            self.push(f"return children")
    
    def resolve_OneOrMore(self, item: OneOrMore, c, queue):
        self.push(f"def _loop_{c}(self):")
        with self.indent():
            self.push(f'"""\n{GrammarPrinter(None).process(item)}\n"""')
            self.push("children = []")
            self.push("while True:")
            with self.indent():
                self.push("pos = self.mark()")
                self.gen(item.expr, queue)
                self.push("if self.match(part): children.append(part)")
                self.push("else:")
                with self.indent():
                    self.push("if len(children) == 0:")
                    with self.indent():
                        self.push("self.fail()")
                    self.push("self.goto(pos)")
                    self.push("break")
            self.push("return children if len(children) > 0 else None")
    
    def resolve_ZeroOrOne(self, item: ZeroOrOne, c, queue):
        self.push(f"def _maybe_{c}(self):")
        with self.indent():
            self.push(f'"""\n{GrammarPrinter(None).process(item)}\n"""')
            self.push("pos = self.mark()")
            self.gen(item.expr, queue)
            self.push("if self.match(part): return part")
            self.push("self.goto(pos)")
            self.push("return Filler()")
    
    def resolve_ExprList(self, item: ExprList, c, queue):
        self.push(f"def _expr_list_{c}(self):")
        with self.indent():
            self.push(f'"""\n{GrammarPrinter(None).process(item)}\n"""')
            self.push("pos = self.mark()")
            self.push("parts = []")
            self.push("for _ in range(1):")
            with self.indent():
                for part in item.exprs:
                    self.gen(part, queue)
                    if not isinstance(part, Predicate):
                        # don't generate extra clause for Predicates
                        self.push("if not self.match(part):")
                        with self.indent():
                            self.push("self.fail()")
                            self.push("break")
                        self.push("parts.append(part)")     
                self.push("return parts")
            self.push("self.goto(pos)")
            self.push("return None")
    
    def resolve_OrOp(self, item: OrOp, c, queue):
        self.push(f"def _or_{c}(self):")
        with self.indent():
            self.push(f'"""\n{GrammarPrinter(None).process(item)}\n"""')
            self.push("pos = self.mark()")
            for choice in item.exprs:
                self.gen(choice, queue)
                self.push("if self.match(part): return part")
                self.push("self.goto(pos)")
            self.push("self.fail()")
            self.push("return None")

    def resolve_queue(self, queue):
        for item, c in queue:
            q = []
            self.resolve(item, c, q)
            if len(q) > 0:
                self.resolve_queue(q)


