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 ConstantsToVars(Pass):
    """
    Ensures every `Update` arg list contains only `Var` nodes, never `Literal`s.
    1. Replaces the literal with a fresh `Var` named `<type>_<index>`.
    2. Emits an equality `eq(new_var, literal)` into the enclosing `Logical`.
    """
    def _rewrite_task(self, node: mm.Task, config:Config) -> mm.Task:
        r = self.ConstantToVarRewriter()
        return r(node)

    class ConstantToVarRewriter(Rewriter):
        def logical(self, node: mm.Logical) -> mm.Logical:
            # In order to recurse over subtasks.
            changed = False

            new_body = []
            for subtask in node.body:
                if isinstance(subtask, mm.Update):
                    modified_subtasks = deconst_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

def replace_constants_with_vars(vals: Sequence[mm.Value]) -> Tuple[List[mm.Value], List[mm.Lookup]] | None:
    new_vals = []
    eqs = []
    changed = False

    for i, val in enumerate(vals):
        if isinstance(val, mm.Literal):
            changed = True
            # Replace constant with a new Var.
            typ = val.type
            assert isinstance(typ, mm.ScalarType), "can only replace scalar constants with vars"
            new_var = mm.Var(typ, f"{typ.name.lower()}_{i}")
            new_vals.append(new_var)
            eqs.append(mm.Lookup(builtins.core.eq, (new_var, val)))
        else:
            new_vals.append(val)

    if not changed:
        return None
    return new_vals, eqs

def deconst_update(update: mm.Update) -> List[Union[mm.Update, mm.Lookup]] | None:
    result = replace_constants_with_vars(update.args)
    if result is None:
        return None
    deduped_vals, req_lookups = result
    new_update = update.replace(args=tuple(deduped_vals))
    return req_lookups + [new_update]
