from ...metamodel.rewriter import Rewriter
from ...metamodel import metamodel as mm
from ...experimental.loopy.builtins import mk_assign
from relationalai.util.structures import OrderedSet

#--------------------------------------------------
# RewriterWithIntermediates
#--------------------------------------------------
# Extends the base Rewriter with the generic ability to add extra tasks to the current
# model. This rewriter also correctly handles adding extra tasks inside Sequence nodes
# preserving the correct order.

class RewriterWithIntermediates(Rewriter):
    def __init__(self, root_logical: mm.Logical):
        super().__init__()
        self.root_logical_id: int = root_logical.id
        # Track tasks that should be added to the top-level
        self.intermediate_tasks: OrderedSet[mm.Task] = OrderedSet()
        # Stack for save/restore of intermediate_tasks around Sequence boundaries
        self._saved_intermediate_tasks: list[OrderedSet[mm.Task]] = []

    # Add a task computing an intermediate relation. This rewriter will automatically handle
    # inserting it at the right level (top-level or inside Sequence).
    def add_intermediate_task(self, task: mm.Task):
        self.intermediate_tasks.add(task)

    def logical(self, node: mm.Logical) -> mm.Logical:
        if node.id != self.root_logical_id:
            # Only add intermediate tasks at the root level, not inside nested Logicals
            return node

        if not self.intermediate_tasks:
            return node

        new_body = tuple(self.intermediate_tasks) + (node,)
        self.intermediate_tasks = OrderedSet()
        return mm.Logical(body=new_body, reasoner=node.reasoner)

    def enter_sequence(self, node: mm.Sequence):
        # Save intermediate_tasks from outside this Sequence
        self._saved_intermediate_tasks.append(self.intermediate_tasks)
        self.intermediate_tasks = OrderedSet()

        # Manually rewrite each task, inserting intermediate tasks
        # right before the task that generated them. This preserves ordering so
        # that intermediates don't run before the Sequence steps they depend on.
        interleaved_sequence_body: list[mm.Task] = []
        for task in node.tasks:
            rewritten = self._rewrite(task)
            for intermediate_task in self.intermediate_tasks:
                # Inside of a Sequence, new defs need to be marked `@assign` to actually
                # define a new relation
                new_assign_task = mk_assign(intermediate_task)
                assert isinstance(new_assign_task, mm.Task)
                interleaved_sequence_body.append(new_assign_task)

            self.intermediate_tasks = OrderedSet()
            assert isinstance(rewritten, mm.Task)
            interleaved_sequence_body.append(rewritten)

        self.intermediate_tasks = OrderedSet(interleaved_sequence_body)

        # Return node with empty tasks so the generated rewriter
        # doesn't re-rewrite children (they've already been rewritten).
        return node.mut(tasks=())

    def sequence(self, node: mm.Sequence) -> mm.Sequence:
        result = node.mut(tasks=tuple(self.intermediate_tasks))

        # Restore intermediate_tasks from outside this Sequence
        self.intermediate_tasks = self._saved_intermediate_tasks.pop()

        return result
