from typing import cast

from . import Pass
from ....metamodel import metamodel as mm
from ....metamodel.rewriter import Walker
from ..rewriter import RewriterWithIntermediates
from .....config import Config
from ..dependency import analyze
from .utils import NameGenerator, create_connection_relation, negate_lookup, compute_exposed_vars


class Match2Union(Pass):
    """
    Handles Match nodes by extracting each branch into a top-level logical under a Model
    or Sequence.

    Each branch writes into its own temporary relation, and the match is replaced
    by a union of lookups to these temporary relations. Each branch (after the first)
    also negates all previous branches.

    From:
        Logical
            Logical
                Match
                    Logical
                        lookup1(...)
                    Logical
                        lookup2(...)
                output(...)
    To:
        Logical
            Logical{…}
                lookup1(...)
                → derive tmp1(...)
            Logical{…}
                lookup2(...)
                Not
                    tmp1(...)
                → derive tmp2(...)
            Logical
                Union
                    tmp1(...)
                    tmp2(...)
                output(...)
    """

    def __init__(self):
        super().__init__()
        self.name_gen = NameGenerator()

    def reset(self):
        self.name_gen.reset()

    def _rewrite_task(self, node: mm.Task, config: Config) -> mm.Task:
        dep_info = analyze(node)

        if not isinstance(node, mm.Logical):
            return node

        # Gather info - each Match node maps to the tasks that need to be added for each branch and the Union reference to the branches
        info_getter = MatchInfoGetter(dep_info)
        info_getter(node)
        match_info = info_getter.match_info

        # Perform the rewrite
        rewriter = MatchFlattener(match_info, node)
        new_root = cast(mm.Task, rewriter(node))
        return new_root

class MatchInfoGetter(Walker):
    """
    Gathers information needed to handle Match nodes.

    For each Match node, collects:
    1. A list of Logical tasks that compute temporary relations for each branch
       (including dependencies and negations of earlier branches)
    2. A Union (or single Lookup) that references all branch relations to replace
       the original Match node
    """
    def __init__(self, dep_info):
        super().__init__()
        self.dep_info = dep_info
        self.name_gen = NameGenerator()

        # Keep a mapping of Match nodes to their generated branch relations and the Union
        # references to those branch relations
        self.match_info: dict[mm.Match, tuple[list[mm.Task], mm.Task]] = {}

    def enter_match(self, node: mm.Match):
        # TODO: handle the incorrect case of an empty Match body more gracefully
        assert node.tasks, "Match node must have at least one branch"

        # Get the exposed variables for this match
        exposed_vars = compute_exposed_vars(node, self.dep_info)

        # For this Match node, we need to record:
        #   1. The Logical tasks that build up the branch relation for each branch
        #   2. The Union reference to these branch relations that will replace the Match node
        branch_relation_tasks: list[mm.Task] = []

        # If the Match has a single branch, then just flatten it into its singular body.
        if len(node.tasks) == 1:
            self.match_info[node] = (branch_relation_tasks, node.tasks[0])
            return self

        # Get dependencies for the match (tasks that must execute before any branch)
        deps = self.dep_info.task_dependencies(node)

        # Determine how many wildcards to use when negating references
        outputs = self.dep_info.task_outputs(node)
        negation_len = len(outputs) if outputs else 0

        # References to each branch's connection relation
        branch_relation_references = []

        # The negated reference to add for the next branch
        negated_references: list[mm.Not] = []

        # Process each branch
        for branch in node.tasks:
            # Create connection relation for this branch
            name = self.name_gen.next("match")
            connection = create_connection_relation(exposed_vars, name)

            # Build the extracted logical body
            branch_body: list[mm.Task] = list(deps)

            # Add the rewritten branch to the body
            if isinstance(branch, mm.Logical):
                branch_body.extend(branch.body)
            else:
                branch_body.append(branch)

            # Add extra tasks (negations + derive)
            branch_body.extend(negated_references)

            branch_update = mm.Update(
                relation=connection,
                args=tuple(exposed_vars),
                effect=mm.Effect.derive
            )
            branch_body.append(branch_update)

            # Extract as a top-level logical
            branch_relation_task = mm.Logical(
                body=tuple(branch_body),
                scope=True,
                reasoner=node.reasoner
            )

            # Keep track of the task which computes the relation for this branch
            branch_relation_tasks.append(branch_relation_task)

            # Create reference to this branch
            reference = mm.Lookup(
                relation=connection,
                args=tuple(exposed_vars)
            )
            branch_relation_references.append(reference)

            # Create negated reference for next branch
            negated_references.append(negate_lookup(reference, negation_len))

        # Return a Union of all branch references
        assert len(branch_relation_references) > 1, "should be more than one branch if we've reached this point"
        branches_reference = mm.Union(
            tasks=tuple(branch_relation_references),
            outputs=node.outputs
        )

        self.match_info[node] = (branch_relation_tasks, branches_reference)
        return self

class MatchFlattener(RewriterWithIntermediates):
    """
    Rewriter that handles Match nodes.

    Replaces each Match node with a Union (or Lookup) that references the temporary
    relations computed by each branch. The Logical tasks that compute these temporary
    relations are collected in extra_tasks and added at the root level.

    Special case: When a Match appears inside a Sequence node, the extra tasks must be
    added at the Sequence level (not the root) and wrapped with @assign annotations.
    """

    def __init__(self, match_info, root_logical: mm.Logical):
        super().__init__(root_logical)
        self.match_info: dict[mm.Match, tuple[list[mm.Task], mm.Task]] = match_info

    def match(self, node: mm.Match) -> mm.Task:
        assert node in self.match_info

        branch_relation_tasks, branches_reference = self.match_info[node]

        # Each branch relation task may contain more Match nodes that need to be flattened,
        # so rewrite them first
        for branch_relation_task in branch_relation_tasks:
            rewritten_task = self._rewrite(branch_relation_task)
            assert isinstance(rewritten_task, mm.Task)
            self.add_intermediate_task(rewritten_task)

        return branches_reference
