import json
from copy import deepcopy
from collections import defaultdict

from tatsu.model import NodeWalker
from typing import List, Dict

from textworld.utils import check_flag
from textworld.logic import Rule
from textworld.envs.pddl.textgen.model import CSGModelBuilderSemantics
from textworld.envs.pddl.textgen.parser import CSGParser


class Symbol:
    def __init__(self, symbol: str, context: Dict = {}):
        self.symbol = symbol
        self.context = context

    def __str__(self):
        return str(self.symbol)

    def copy(self, context):
        copy = deepcopy(self)
        copy.context = copy_context(context)
        return copy


class TerminalSymbol(Symbol):

    def __repr__(self):
        return "TerminalSymbol({})".format(repr(self.symbol))


class NonterminalSymbol(Symbol):

    def __repr__(self):
        return "NonterminalSymbol('#{}#')".format(self.symbol)


def copy_context(context):
    return {
        "state": context["state"],
        "facts": context["facts"],
        "variables": deepcopy(context["variables"]),
        "mapping": deepcopy(context["mapping"]),
        "entity_infos": context["entity_infos"],
    }


def join(sep, iterable):
    """ Join an iterable with a separator similar to `str.join`. """
    iterable = iter(iterable)
    yield next(iterable)

    for e in iterable:
        yield sep
        yield e


def display_list(l, context):
    if len(l) == 0:
        return [NonterminalSymbol("list_empty", context)]

    if len(l) == 1:
        return [l[0]]

    list_separator = NonterminalSymbol("list_separator", context)
    list_last_separator = NonterminalSymbol("list_last_separator", context)
    return list(join(list_separator, l[:-1])) + [list_last_separator] + [l[-1]]


def query(expression, context):
    rule = Rule.parse_conjunctive_query(expression)

    contexts = []
    for mapping in context["state"].all_assignments(rule, context["mapping"]):
        context_ = copy_context(context)
        new_variables = {ph.name: context_["entity_infos"][var.name] for ph, var in mapping.items()}
        context_["variables"].update(new_variables)
        context_["mapping"].update(mapping)
        contexts.append(context_)

    return contexts


class EvalSymbol(Symbol):
    def __init__(self, expression: str, context: Dict = {}):
        super().__init__(expression, context)
        self.expression = expression

    def __repr__(self):
        return "EvalSymbol('{{{}}}')".format(str(self.expression))

    def derive(self, context=None):
        context = context or self.context
        locals().update(context["variables"])
        value = eval(self.expression)
        return [TerminalSymbol(value)]


class ListSymbol(Symbol):
    def __init__(self, symbol: Symbol, context: Dict = {}):
        super().__init__(symbol, context)

    def __repr__(self):
        return "ListSymbol('[{!r}]')".format(self.symbol)

    def derive(self, context=None):
        context = context or self.context
        self.symbol.context = context
        derivation = self.symbol.derive()
        return display_list(derivation, context)


class ConditionalSymbol(Symbol):

    def __init__(self, expression: Symbol, given: str, context: Dict = {}):
        super().__init__(str(expression), context)
        self.expression = expression
        self.given = given

    def __repr__(self):
        return "ConditionalSymbol('{{{}|{}}}')".format(str(self.expression), str(self.given))

    def derive(self):
        if len(self.context) == 0:
            raise ValueError("Empty context")

        context = copy_context(self.context)

        contexts = [context]
        if self.given:
            contexts = query(self.given, context)

        res = [self.expression.copy(context) for context in contexts]
        return res


class ProductionRule:
    """ Production rule for a context-sensitive grammar. """

    def __init__(self, lhs: str, rhs: List[Symbol], weight=1, condition=None):
        """
        Arguments:
            rhs: symbol that will be transformed by this production rule.
            lhs: list of symbols generated by this production rule.
            weight: prevalence of this production.
        """
        self.lhs = lhs
        self.rhs = rhs
        self.weight = weight
        self.condition = condition

    def __repr__(self):
        text = "ProductionRule(lhs={!r}, rhs={!r}, weight={!r}, condition={!r})"
        return text.format(self.lhs, self.rhs, self.weight, self.condition)


class _Converter(NodeWalker):

    def walk_list(self, node):
        return [self.walk(child) for child in node]

    def walk_str(self, node):
        return node.replace("\\n", "\n")

    def walk_TerminalSymbol(self, node):
        return TerminalSymbol(node.literal)

    def walk_NonterminalSymbol(self, node):
        return NonterminalSymbol(node.symbol)

    def walk_ConditionalSymbol(self, node):
        return ConditionalSymbol(self.walk(node.expression), node.given)

    def walk_EvalSymbol(self, node):
        return EvalSymbol(node.statement)

    def walk_ListSymbol(self, node):
        return ListSymbol(self.walk(node.symbol))


_PARSER = CSGParser(semantics=CSGModelBuilderSemantics(), parseinfo=True)


def _parse_and_convert(*args, **kwargs):
    model = _PARSER.parse(*args, **kwargs)
    return _Converter().walk(model)


class CSGUnknownSymbolError(Exception):
    def __init__(self, symbol: Symbol):
        msg = "Can't find symbol '#{}#' in the set of production rules."
        # TODO: mention closest match?
        super().__init__(msg.format(symbol))


class ContextSensitiveGrammar:

    def __init__(self):
        self._rules = defaultdict(list)

    def update(self, grammar: "ContextSensitiveGrammar"):
        for k, v in grammar._rules.items():
            self._rules[k].extend(v)

    @classmethod
    def parse(cls, text: str):
        data = json.loads(text)

        grammar = cls()
        for name, rules in data.items():
            rules = [ProductionRule(lhs=name,
                                    rhs=_parse_and_convert(rule["rhs"], rule_name="symbols"),
                                    weight=rule.get("weight", 1),
                                    condition=rule.get("condition", ""))
                     for rule in rules]

            for rule in rules:
                grammar.add_rule(rule)

        return grammar

    def add_rule(self, rule: ProductionRule):
        self._rules[rule.lhs].append(rule)

    def replace(self, start: Symbol) -> List[Symbol]:
        rules = self._rules.get(str(start))
        if not rules:
            raise CSGUnknownSymbolError(start)

        def _applicable(rule):
            if not rule.condition:
                return True

            rule = Rule.parse_conjunctive_query(rule.condition)
            mappings = list(start.context["state"].all_assignments(rule, start.context["mapping"]))
            return len(mappings) > 0  # Found a match.

        rules = list(filter(_applicable, rules))

        # TODO: deal with multiple alternatives
        # TODO: deal with context
        symbols = deepcopy(rules[0].rhs)
        for symbol in symbols:
            symbol.context = copy_context(start.context)

        return symbols

    def derive(self, start: str, context={}) -> str:
        derivation = _parse_and_convert(start, rule_name="symbols", trace=check_flag("TW_CSG_TRACE"))
        derivation = derivation[::-1]  # Reverse to build a derivation stack.

        for symbol in derivation:
            symbol.context = copy_context(context)

        derived = []
        while len(derivation) > 0:
            if check_flag("TW_CSG_DEBUG"):
                print(derivation)

            symbol = derivation.pop()
            if isinstance(symbol, TerminalSymbol):
                derived.append(symbol)

            elif isinstance(symbol, NonterminalSymbol):
                derivation += self.replace(symbol)[::-1]  # Reverse to add on top of the derivation stack.

            elif isinstance(symbol, (ConditionalSymbol, EvalSymbol, ListSymbol)):
                derivation += symbol.derive()[::-1]  # Reverse to add on top of the derivation stack.

            else:
                raise NotImplementedError("Unknown symbol: {}".format(type(symbol)))

        return "".join(map(str, derived))
