from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, TypeVar
from typing import Tuple

from relationalai.semantics.metamodel import builtins, ir, factory as f, helpers
from relationalai.semantics.metamodel.compiler import Pass, group_tasks
from relationalai.semantics.metamodel.util import NameCache, OrderedSet, ordered_set
from relationalai.semantics.metamodel import dependency
from relationalai.semantics.metamodel.typer.typer import to_type

from ..builtins import lifted_data

class Flatten(Pass):
    """
    Traverses the model's root to flatten it as much as possible. The result of this pass is
    a Logical root where all nested tasks that represent a rule in Rel are extracted to the
    top level. Additionally, any Sequence is promoted to the top level Logical (but
    encapsulated by a Logical).

    - nested logical with updates becomes a top-level logical (a rule)

    From:
        Logical
            Logical
                lookup1   <- scope is spread
                Logical
                    lookup2
                    derive foo
                Logical
                    lookup3
                    derive bar
    To:
        Logical
            Logical
                lookup1
                lookup2
                derive foo
            Logical
                lookup1
                lookup3
                derive bar

    - a match becomes a top-level logical for each branch, each writing into its own temporary
    relation and a lookup from the last relation. The top-level logical for a branch derives
    into the temporary relation negating the previous branch:

    From:
        Logical
            Logical
                Match
                    Logical
                        lookup1
                    Logical
                        lookup2
                output
    To:
        Logical
            Logical
                lookup1
                derive tmp1
            Logical
                Union            <- tmp1() or (not temp1() and lookup2())
                    lookup tmp1
                    Logical
                        Not
                            lookup tmp1
                        lookup2
                        derive tmp2
            Logical
                lookup tmp2
                output

    - a Sequence is promoted to the top level Logical, encapsulated by a Logical:
    From:
        Logical
            Logical
                lookup
                derive foo
            Sequence
                Logical
                    ...
                Loop
                    Sequence
                        ...
                Logical
                    ...
    To:
        Logical
            Logical
                lookup
                derive foo
            Logical
                Sequence
                    Logical
                        ...
                    Loop
                        Sequence
                            ...
                    Logical
                        ...
    """

    def __init__(self):
        super().__init__()
        self.name_cache = NameCache(start_from_one=True)

    #--------------------------------------------------
    # Public API
    #--------------------------------------------------
    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
        # create the context
        ctx = Flatten.Context(model, options)

        # rewrite the root
        result = self.handle(model.root, ctx)

        # the new body contains the extracted top level logicals and maybe the rewritten root
        body = ctx.rewrite_ctx.top_level if result.replacement is None else ctx.rewrite_ctx.top_level + [result.replacement]

        # create the new model, updating relations and root
        return ir.Model(
            model.engines,
            OrderedSet.from_iterable(model.relations).update(ctx.rewrite_ctx.relations).frozen(),
            model.types,
            ir.Logical(model.root.engine, tuple(), tuple(body))
        )

    #--------------------------------------------------
    # IR handlers
    #--------------------------------------------------

    class Context():
        def __init__(self, model: ir.Model, options: dict):
            self.rewrite_ctx = helpers.RewriteContext()
            self.info = dependency.analyze_constraints(model.root)
            self.options = options
            # when we extract a nested logical (due to match, etc), we also
            # extract together its dependencies, and we leave behind a different nested
            # logical as a reference, with a lookup to the extracted relation. This dict
            # keeps a map from the reference logical id to the set of tasks that where
            # extracted for it. This is used to optimize the resulting code, since we can
            # assume that the reference logical already includes those tasks.
            self.included: dict[int, OrderedSet[ir.Task]] = dict()
            # These are extra tasks to be added to the body of a logical when it is
            # extracted as a top-level logical. For example, when handling a match, we add
            # a couple of tasks here in case the branches have updates or aggregations that
            # cause the branch to be extracted, so that they include these tasks.
            self.extra_tasks = []


    @dataclass
    class HandleResult():
        replacement: Optional[ir.Task]

    def handle(self, task: ir.Task, ctx: Context) -> Flatten.HandleResult:
        if isinstance(task, ir.Logical):
            return self.handle_logical(task, ctx)
        elif isinstance(task, ir.Match):
            return self.handle_match(task, ctx)
        elif isinstance(task, ir.Require):
            return self.handle_require(task, ctx)
        elif isinstance(task, ir.Sequence):
            return self.handle_sequence(task, ctx)
        else:
            return Flatten.HandleResult(task)


    def handle_logical(self, task: ir.Logical, ctx: Context):
        # keep track of what's the result of handling nested composites
        composites = group_tasks(task.body, {
            "composites": helpers.COMPOSITES
        })["composites"]
        all_composites_removed = len(composites) > 0

        if any(anno.relation is lifted_data for anno in task.annotations):
            # if this is a lifted data logical, we can skip handling of the body and
            # flatten it directly.
            ctx.rewrite_ctx.top_level.append(task)
            return Flatten.HandleResult(None)

        # recursively handle children, collecting the replacements in the body
        body:OrderedSet[ir.Task] = ordered_set()
        for child in task.body:
            result = self.handle(child, ctx)
            if result.replacement is not None:
                ctx.info.replaced(child, result.replacement)
                extend_body(body, result.replacement)
                # nested composite was not completely removed
                if child in composites:
                    all_composites_removed = False

        # Filter out empty logicals from the body
        new_body:OrderedSet[ir.Task] = ordered_set()
        for b in body:
            if isinstance(b, ir.Logical):
                if not b.body:
                    # empty logical, skip
                    continue
            new_body.add(b)

        body = new_body

        # all children were extracted or all composites were removed without any effects
        # left and no outputs (so no way for outer dependencies), drop this logical
        if not body or (all_composites_removed and not any([isinstance(t, helpers.EFFECTS) for t in body]) and not ctx.info.task_outputs(task)):
            return Flatten.HandleResult(None)

        # now process the rewritten body
        groups = group_tasks(body.get_list(), {
            "outputs": ir.Output,
            "updates": ir.Update,
        })

        # If there are outputs, flatten each into its own top-level rule, along with its
        # dependencies.
        if groups["outputs"]:
            # Analyze the dependencies in the newly rewritten body
            new_logical = ir.Logical(task.engine, task.hoisted, tuple(body))
            info = dependency.analyze(new_logical)

            for output in groups["outputs"]:
                assert(isinstance(output, ir.Output))
                deps = info.task_dependencies(output)
                new_body = OrderedSet.from_iterable(t.clone() for t in deps)
                new_body.update(t.clone() for t in ctx.extra_tasks)
                new_body.add(output.clone())
                ctx.rewrite_ctx.top_level.append(ir.Logical(task.engine, task.hoisted, tuple(new_body), task.annotations))

            return Flatten.HandleResult(None)

        # if there are updates, extract as a new top level rule
        if groups["updates"]:
            # add task dependencies to the body
            body.prefix(t.clone() for t in ctx.info.task_dependencies(task))
            # potentially add context extra tasks
            body.update(t.clone() for t in ctx.extra_tasks)
            ctx.rewrite_ctx.top_level.append(ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations))
            return Flatten.HandleResult(None)

        return Flatten.HandleResult(ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations))


    def handle_match(self, match: ir.Match, ctx: Context):
        # TODO: how to deal with malformed input like this?
        if not match.tasks:
            return Flatten.HandleResult(match)

        body = ctx.info.task_dependencies(match)
        exposed_vars = self.compute_exposed_vars(match, match.tasks, ctx)

        # negation len is the number of wildcards when negating a reference
        outputs = ctx.info.task_outputs(match)
        negation_len = len(outputs) if outputs else 0

        # collect the references to branch rules so that we union in the end
        references = []

        # a negated reference to the previous branch (e.g. "not branch1(...)")
        negated_reference = None

        # number of negated references currently in context (e.g. in the 3rd branch there
        # will be 2 negated references, "not branch1(...) and not branch2(...)")
        negated_references = 0
        for branch in match.tasks:
            # if there's a previous connection, add the negation of it to extra tasks
            if negated_reference:
                ctx.extra_tasks.append(negated_reference)
                negated_references += 1

            # the connection relation to be used for this branch
            name = helpers.create_task_name(self.name_cache, branch, "_match")
            relation = helpers.create_connection_relation(branch, exposed_vars, ctx.rewrite_ctx, name)
            branch_update = f.derive(relation, exposed_vars)
            ctx.extra_tasks.append(branch_update)

            # handle the branch
            result = self.handle(branch, ctx)

            # if there's no replacement, the branch was already extracted (e.g. by an update),
            # otherwise, we need to extract the branch as a rule writing into the connection
            if result.replacement:
                ctx.info.replaced(branch, result.replacement)

                branch_body: OrderedSet[ir.Task] = OrderedSet.from_iterable(body)
                extend_body(branch_body, result.replacement)
                # potentially add context extra tasks (negations of previous branches + derive into this branch)
                branch_body.update(ctx.extra_tasks)
                # extract the body
                ctx.rewrite_ctx.top_level.append(ir.Logical(match.engine, tuple(), tuple(branch_body)))

            # cleanup the branch update from context (but keep the negations)
            ctx.extra_tasks.pop()

            # record a reference to the branch connection and adjust the negated_reference
            reference = f.lookup(relation, exposed_vars)
            negated_reference = negate(reference, negation_len)
            references.append(reference)

        # we accumulate negations, in the end we remove them all
        for x in range(negated_references):
            ctx.extra_tasks.pop()

        # return a union of the references to the branches
        return Flatten.HandleResult(f.union(references, match.hoisted))

    def compute_exposed_vars(self, task: ir.Task, branches: Tuple[ir.Task, ...], ctx: Context) -> list:
        """ Helper to compute the set of exposed vars for handle_match and handle_union. """
        exposed_vars = set_union(ctx.info.task_inputs(task), ctx.info.task_outputs(task))
        if exposed_vars:
            return exposed_vars
        # no exposed vars, so we try to use the common vars across all branches as context
        common_vars = helpers.collect_vars(branches[0])
        for branch in branches:
            branch_vars = helpers.collect_vars(branch)
            for v in common_vars:
                if v not in branch_vars:
                    common_vars.remove(v)
        # TODO - we should raise an error here if the set is empty
        return common_vars.get_list()


    def handle_require(self, req: ir.Require, ctx: Context):
        # only extract the domain if it is a somewhat complex Logical and there's more than
        # one check, otherwise insert it straight into all checks
        if builtins.discharged_annotation in req.annotations:
            # remove discharged Requires
            return Flatten.HandleResult(None)
        elif builtins.declare_constraint_annotation in req.annotations:
            # leave Requires that are declared constraints
            return Flatten.HandleResult(req)
        else:
            # generate logic for remaining requires
            domain = req.domain
            if len(req.checks) > 1 and isinstance(domain, ir.Logical) and len(domain.body) > 1:
                body = OrderedSet.from_iterable(domain.body)
                vars = helpers.hoisted_vars(domain.hoisted)
                name = helpers.create_task_name(self.name_cache, req)
                connection = helpers.extract(req, body, vars, ctx.rewrite_ctx, name)
                domain = f.logical([f.lookup(connection, vars)], vars)

            for check in req.checks:
                # only generate logic for checks that have errors and not discharged
                if check.error and builtins.discharged_annotation not in check.annotations:
                    handled_check_result = self.handle(check.check, ctx)
                    if handled_check_result.replacement:
                        body = ordered_set()
                        body.add(domain)
                        body.add(ir.Not(req.engine, handled_check_result.replacement))
                        if (isinstance(check.error, ir.Logical)):
                            body.update(check.error.body)
                        else:
                            # this is more general but may trip the current splinter
                            body.add(check.error)
                        ctx.rewrite_ctx.top_level.append(ir.Logical(req.engine, tuple(), tuple(body)))

        # currently we just drop the Require, but we should keep it here and link the
        # extracted logicals to it
        return Flatten.HandleResult(None)

    def handle_sequence(self, task: ir.Sequence, ctx: Context):
        new_logical = f.logical(
            body = [task],
            engine = task.engine
        )
        ctx.rewrite_ctx.top_level.append(new_logical)
        return Flatten.HandleResult(None)

#--------------------------------------------------
# Helpers
#--------------------------------------------------

T = TypeVar('T')
def set_union(s1: Optional[OrderedSet[T]], s2: Optional[OrderedSet[T]]) -> list:
    """ Return the union of s1 | s2, accounting for when any is None. """
    if s1 and s2:
        return (s1 | s2).get_list()
    if s1:
        return s1.get_list()
    if s2:
        return s2.get_list()
    return []

def negate(lookup: ir.Lookup, values: int):
    """
    Return a negation of this reference, where the last `values` arguments are to
    be replaced by wildcards (i.e. len(reference.args) - values are keys so they need
    to be bound in the Not.)
    """
    args = []
    i = 0
    last = len(lookup.args) - values
    for arg in lookup.args:
        args.append(f.wild(to_type(arg))) if i >= last else args.append(arg)
        i += 1

    return ir.Not(lookup.engine, f.lookup(lookup.relation, args))

def merge_var_list(vars: list[ir.Var], hoisted: tuple[ir.VarOrDefault, ...]) -> list[ir.VarOrDefault]:
    """ Merge vars and hoisted, making sure that hoisted vars have precedence since they may have defaults. """
    r = []
    hoisted_vars = helpers.hoisted_vars(hoisted)
    for v in vars:
        if v not in hoisted_vars:
            r.append(v)
    r.extend(hoisted)
    return r

def extend_body(body: OrderedSet[ir.Task], extra: ir.Task):
    """ Add the extra task to the body, but if the extra is a simple logical, just
    inline its subtasks. """
    if isinstance(extra, ir.Logical):
        if extra.hoisted:
            # hoists, remove things that are already in the body to avoid duplicates
            logical_body = []
            for t in extra.body:
                if t not in body:
                    logical_body.append(t)
            if len(logical_body) == len(extra.body):
                # no duplicates
                body.add(extra)
            else:
                # some duplicate, remove them
                body.add(ir.Logical(
                    extra.engine,
                    extra.hoisted,
                    tuple(logical_body)
                ))
        else:
            body.add(extra)
    else:
        body.add(extra)
