from __future__ import annotations

from relationalai.semantics.metamodel import ir, factory as f, helpers
from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
from relationalai.semantics.lqp.rewrite.flatten import Flatten, negate, extend_body
from relationalai.semantics.lqp.algorithms import is_script, mk_assign

class FlattenScript(Flatten):
    """
    Flattens Match nodes inside @script Sequence blocks. This pass extends Flatten
    of standard Logicals, to reuse a number of utilities, but DOES NOT flatten Match
    nodes outside of scripts (which are handled by the Flatten pass).

    Unlike the regular Flatten pass which extracts to top-level, FlattenScript
    maintains order by inserting intermediate relations right before they're used
    within the Sequence. This is necessary because order matters in a script Sequence.

    Additionally, if the original Logical is a Loopy instruction, the flattened Logicals are
    made Loopy too (`@assign`).

    Example:

    === BEFORE ===
    Logical
        Sequence @script @algorithm
            Logical
                dom(n)
                Match ⇑[k]
                    Logical ⇑[k]
                        value(n, k)
                    Logical ⇑[k]
                        k = 0
                → derive result(n, k) @assign @global
                filter(n)

    === AFTER ===
    Logical
        Sequence @script @algorithm
            Logical
                dom(n)
                filter(n)
                Logical ⇑[v]
                    value(n, v)
                → derive _match_1(n, v) @assign
            Logical
                dom(n)
                filter(n)
                Logical ⇑[v]
                    v = 0
                Not
                    _match_1(n, _)
                → derive _match_2(n, v) @assign
            Logical
                dom(n)
                filter(n)
                Union ⇑[v]
                    _match_1(n, v)
                    _match_2(n, v)
                → derive result(n, v) @assign @global
    """

    class Context(Flatten.Context):
        """Extended context with script tracking."""
        def __init__(self, model: ir.Model, options: dict):
            super().__init__(model, options)
            self.in_script: bool = False

    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
        """Traverse the model and flatten Match nodes inside script Sequences."""
        ctx = FlattenScript.Context(model, options)
        result = self.handle(model.root, ctx)

        if result.replacement is None:
            return model

        # Convert relations list to FrozenOrderedSet (adding any new intermediate relations)
        new_relations = OrderedSet.from_iterable(model.relations).update(ctx.rewrite_ctx.relations).frozen()

        return ir.Model(
            model.engines,
            new_relations,
            model.types,
            result.replacement
        )

    def handle(self, task: ir.Task, ctx: Flatten.Context) -> Flatten.HandleResult:
        """Override handle to add Loop support."""
        if isinstance(task, ir.Loop):
            return self.handle_loop(task, ctx)
        return super().handle(task, ctx)

    def handle_loop(self, task: ir.Loop, ctx: Flatten.Context) -> Flatten.HandleResult:
        """Recursively handle the body of the loop."""
        result = self.handle(task.body, ctx)

        assert(result.replacement)

        # If body unchanged, return original loop
        if result.replacement is task.body:
            return Flatten.HandleResult(task)

        # Return new loop with handled body
        return Flatten.HandleResult(ir.Loop(
            task.engine,
            task.hoisted,
            task.iter,
            result.replacement,
            task.concurrency,
            task.annotations
        ))

    def handle_logical(self, task: ir.Logical, ctx: Context):  # type: ignore[override]
        """
        Handle Logical nodes.

        Outside scripts: simple traversal to find nested scripts.
        Inside scripts: prevent extraction to top-level and keep everything in sequence.

        Note: Type checker complains about parameter type narrowing, but this is safe
        because we only create FlattenScript.Context in our own rewrite() method.
        """
        # Recursively process children
        body: OrderedSet[ir.Task] = ordered_set()
        for child in task.body:
            result = self.handle(child, ctx)
            if result.replacement is not None:
                if ctx.in_script and isinstance(result.replacement, ir.Logical) and not result.replacement.hoisted:
                    # Inside script: inline simple logicals without hoisting
                    body.update(result.replacement.body)
                else:
                    body.add(result.replacement)

        if not body:
            return Flatten.HandleResult(None)

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

    def flatten_match_in_logical(self, logical: ir.Logical, match: ir.Match, match_idx: int, ctx: Context) -> list[ir.Logical]:
        """
        Flatten a Match inside a Logical within a script Sequence.
        Returns a list of Logicals to be inserted in sequence.
        """
        if not match.tasks:
            return [logical]

        # Split the logical into: tasks_before, match, tasks_after
        tasks_before = list(logical.body[:match_idx])
        tasks_after = list(logical.body[match_idx + 1:])

        # Separate tasks_after into filters (non-Update tasks) and updates (Update tasks)
        # Filters are constraints that should be included in all branches
        filters = [task for task in tasks_after if not isinstance(task, ir.Update)]
        updates = [task for task in tasks_after if isinstance(task, ir.Update)]

        # Compute exposed variables
        exposed_vars = self.compute_exposed_vars(match, match.tasks, ctx)

        # Use dependency analysis for branch bodies (like flatten.py does)
        branch_dependencies = ctx.info.task_dependencies(match)

        # Collect all dependencies for the final Logical (tasks before + filters after)
        final_dependencies = tasks_before + filters

        # Negation length for wildcards
        outputs = ctx.info.task_outputs(match)
        negation_len = len(outputs) if outputs else 0

        # Result: list of Logicals to insert
        result_logicals = []
        references = []
        negated_reference = None

        # Process each branch
        for branch in match.tasks:
            # Create connection relation 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)

            # Handle the branch (recursively process nested structures)
            result = self.handle(branch, ctx)
            branch_content = result.replacement if result.replacement else branch

            # Update dependency tracking if branch was transformed
            if result.replacement:
                ctx.info.replaced(branch, result.replacement)

            # Build logical for this branch
            branch_body: OrderedSet[ir.Task] = ordered_set()

            # Add dependencies (using dependency analysis, not all tasks)
            branch_body.update(branch_dependencies)

            # Add branch content using extend_body helper
            extend_body(branch_body, branch_content)

            # Add negation of previous branches (after branch content)
            if negated_reference:
                branch_body.add(negated_reference)

            # Add derive to connection relation
            branch_update = f.derive(relation, exposed_vars)
            branch_body.add(branch_update)

            # Create the Logical for this branch
            branch_logical = mk_assign(ir.Logical(match.engine, tuple(), tuple(branch_body)))

            result_logicals.append(branch_logical)

            # Update references for final union
            reference = f.lookup(relation, exposed_vars)
            negated_reference = negate(reference, negation_len)
            references.append(reference)

        # Create final Logical with Union and remaining tasks
        final_body: OrderedSet[ir.Task] = ordered_set()
        final_body.update(final_dependencies)

        # Add union of all branches
        union = f.union(references, match.hoisted)
        final_body.add(union)

        # Add updates that came after the match (filters are already in dependencies)
        final_body.update(updates)

        # Create final logical preserving the original annotations
        final_logical = ir.Logical(logical.engine, logical.hoisted, tuple(final_body), logical.annotations)

        result_logicals.append(final_logical)

        return result_logicals

    def handle_sequence(self, task: ir.Sequence, ctx: Context):  # type: ignore[override]
        """
        Handle a Sequence.

        If it's a script: set context flag and flatten Match nodes within.
        If not a script: simple traversal (Flatten pass already processed it).

        Note: Type checker complains about parameter type narrowing, but this is safe
        because we only create FlattenScript.Context in our own rewrite() method.
        """
        if not is_script(task):
            # Not a script sequence - already processed by Flatten, just return as-is
            return Flatten.HandleResult(task)

        # This is a script - mark context and process with flattening
        old_in_script = ctx.in_script
        ctx.in_script = True

        # Process the sequence tasks
        new_tasks: list[ir.Task] = []
        for child in task.tasks:
            # Check if this child is a Logical with Match that needs flattening
            if isinstance(child, ir.Logical):
                new_tasks.extend(self.try_flatten_logical(child, ctx))
                continue

            # No flattening needed, process normally
            result = self.handle(child, ctx)
            if result.replacement is not None:
                new_tasks.append(result.replacement)

        # Restore context
        ctx.in_script = old_in_script

        return Flatten.HandleResult(
            ir.Sequence(task.engine, task.hoisted, tuple(new_tasks), task.annotations)
        )

    def try_flatten_logical(self, logical: ir.Logical, ctx: Context) -> list[ir.Logical]:
        """
        Flatten all Matches in a Logical.
        Iteratively flattens until no more Matches remain in any of the resulting Logicals.
        """
        worklist = [logical]
        result = []

        while worklist:
            current = worklist.pop()

            # Find first Match in current logical
            match = None
            match_idx = -1
            for i, child in enumerate(current.body):
                if isinstance(child, ir.Match):
                    match = child
                    match_idx = i
                    break

            if match is None:
                # No Match found - this logical is done
                result.append(current)
            else:
                # Flatten and add results back to worklist for further processing.
                # Reverse so that pop() returns them in the original order.
                flattened = self.flatten_match_in_logical(current, match, match_idx, ctx)
                worklist.extend(reversed(flattened))

        return result
