from . import Pass
from ....metamodel import metamodel as mm
from ....metamodel.builtins import builtins
from ....metamodel.rewriter import Rewriter
from .....config import Config

from typing import List, Sequence, Tuple, Union

class DeduplicateVars(Pass):
    """
    Ensures every `Update` arg list contains each `Var` at most once.
    For each duplicate position it:
    1. Replaces the duplicate with a fresh `Var` (suffix `_dup_<index>`).
    2. Emits an equality `eq(new_var, original_var)` into the enclosing `Logical`.
    """
    def _rewrite_task(self, node: mm.Task, config:Config) -> mm.Task:
        r = self.VarDeduplicator()
        return r(node)

    class VarDeduplicator(Rewriter):
        def logical(self, node: mm.Logical):
            new_body = []
            changed = False
            for subtask in node.body:
                if isinstance(subtask, mm.Update):
                    modified_subtasks = dedup_update(subtask)
                    if modified_subtasks is not None:
                        changed = True
                        new_body.extend(modified_subtasks)
                    else:
                        new_body.append(subtask)
                else:
                    new_body.append(subtask)

            return node.replace(body=tuple(new_body)) if changed else node

# Return
# 1) a new list of Values with no duplicates (at the object level) and
# 2) equalities between any original Value and a deduplicated Value.
# or None if there are no duplicates.
def dedup_values(vals: Sequence[mm.Value]) -> Tuple[List[mm.Value], List[mm.Lookup]] | None:
    # If a var is seen more than once, it is a duplicate and we will create
    # a new Var and equate it with the seen one.
    seen_vars = set()

    new_vals = []
    eqs = []

    for i, val in enumerate(vals):
        if isinstance(val, mm.Var):
            if val in seen_vars:
                new_var = mm.Var(val.type, val.name + "_dup_" + str(i))
                new_val = new_var
                new_vals.append(new_val)
                eqs.append(mm.Lookup(builtins.core.eq, (new_var, val)))
            else:
                seen_vars.add(val)
                new_vals.append(val)
        else:
            # No possibility of problematic duplication.
            new_vals.append(val)

    if not seen_vars:
        return None
    return new_vals, eqs

# Returns a replacement update with no duplicate variables, along with any equalities
# needed between any two previously duplicate variables, or None if there are no duplicates.
def dedup_update(update: mm.Update) -> List[Union[mm.Update, mm.Lookup]] | None:
    replacements = dedup_values(update.args)
    if replacements is None:
        return None
    deduped_vals, req_lookups = replacements
    new_update = update.replace(args=tuple(deduped_vals))
    return req_lookups + [new_update]
