"""
Utility functions for LQP rewrite passes.
"""
from ....metamodel import metamodel as mm
from .....util.structures import OrderedSet
from ..dependency import DependencyInfo
from .....shims.helpers import VarFinder


class NameGenerator:
    """Generates unique names for temporary relations."""

    def __init__(self, start_from: int = 1):
        self.counter = start_from - 1

    def next(self, prefix: str = "temp") -> str:
        """Generate the next unique name with the given prefix."""
        self.counter += 1
        return f"_{prefix}_{self.counter}"

    def reset(self):
        """Reset the counter to its initial state."""
        self.counter = 0


def create_connection_relation(vars: list[mm.Var], name: str) -> mm.Relation:
    """
    Create a temporary connection relation with fields corresponding to the given variables.

    Args:
        vars: Variables to use as fields for the relation
        name: Name for the temporary relation

    Returns:
        A new Relation with fields matching the variables
    """
    fields = tuple(mm.Field(name=v.name, type=v.type) for v in vars)
    return mm.Relation(name=name, fields=fields)


def negate_lookup(lookup: mm.Lookup, num_wildcards: int = 0) -> mm.Not:
    """
    Create a negation of a lookup, optionally replacing trailing arguments with wildcards.

    This is useful for negating references in match branches, where we want to negate
    the fact that previous branches matched, but only bind the key variables.

    Args:
        lookup: The lookup to negate
        num_wildcards: Number of trailing arguments to replace with wildcards
                      (0 means negate the lookup as-is)

    Returns:
        A Not task containing the negated lookup

    Example:
        lookup(rel, (x, y, z)) with num_wildcards=1
        becomes: not lookup(rel, (x, y, _))
    """
    if num_wildcards == 0:
        # Simple case: negate the lookup as-is
        return mm.Not(task=lookup)

    args = list(lookup.args)
    last_idx = len(args) - num_wildcards

    for i in range(last_idx, len(args)):
        arg = args[i]
        if isinstance(arg, mm.Var):
            # Create a wildcard with the same type
            args[i] = mm.Var(type=arg.type, name="_")

    negated_lookup = mm.Lookup(
        relation=lookup.relation,
        args=tuple(args)
    )
    return mm.Not(task=negated_lookup)


def compute_exposed_vars(
    task: mm.Task,
    dep_info: DependencyInfo,
) -> list[mm.Var]:
    """
    Compute the exposed variables for a task.

    These are the variables that need to be passed through when extracting
    the task to a top-level logical with a connection relation.

    Args:
        task: The task to analyze
        dep_info: Dependency information for the task tree
        explicit_outputs: Explicit output variables if specified (e.g., from Match.outputs)

    Returns:
        List of variables that should be exposed
    """
    # Get inputs and outputs from dependency analysis
    inputs = dep_info.task_inputs(task)
    outputs = dep_info.task_outputs(task)

    exposed = OrderedSet()
    if inputs:
        exposed.update(inputs)
    if outputs:
        exposed.update(outputs)

    if exposed:
        return list(exposed)

    # If no exposed vars and task is a container with branches, try common vars
    if isinstance(task, (mm.Match, mm.Union)) and task.tasks:
        finder = VarFinder()

        # Collect vars from first branch
        common_vars = finder.find_vars(task.tasks[0])

        # Intersect with vars from other branches
        for branch in task.tasks[1:]:
            branch_vars = finder.find_vars(branch)
            common_vars = common_vars & branch_vars

        return list(common_vars)

    return []
