from comfy_api.latest import io
from .helper_functions import comonLazy, make_zero_like


class MathNodeBase(io.ComfyNode):
    """
    Base class for More Math nodes providing common utilities.
    """

    @classmethod
    def check_lazy_status(cls, **kwargs):
        """
        Generic check_lazy_status implementation.
        Finds the first argument name that isn't a known variable (a-d, w-z) and assumes it's the expression.
        """
        # Known variable inputs (and internal ones we might ignore)
        known_vars = {"a", "b", "c", "d", "w", "x", "y", "z", "pooled_output"}

        # Heuristic: The expression is usually the first argument that is not a known variable.
        # But `comonLazy` takes (expr, a, b, c, d, w, x, y, z).
        # We need to extract the values from kwargs correctly.

        expr_key = None
        # Try finding the expression key
        for k in kwargs:
            if k not in known_vars and isinstance(kwargs[k], str):
                expr_key = k
                break

        # If we can't find it heuristically (e.g. standard naming), fallback or error?
        # Let's try to be safe. If we can't find it, we might be in a node with custom logic like ConditioningMathNode.
        # ConditioningMathNode has 'Tensor' and 'pooled_output' as expressions.

        if expr_key:
            expr = kwargs[expr_key]
            a = kwargs.get("a", None)
            b = kwargs.get("b", [])
            c = kwargs.get("c", [])
            d = kwargs.get("d", [])
            w = kwargs.get("w", 0)
            x = kwargs.get("x", 0)
            y = kwargs.get("y", 0)
            z = kwargs.get("z", 0)
            return comonLazy(expr, a, b, c, d, w, x, y, z)

        # Fallback/Empty return if we can't determine dependencies automatically
        return []

    @staticmethod
    def prepare_inputs(a, b, c, d):
        """
        Ensures optional inputs b, c, d are zero-initialized like a if None.
        Returns the prepared (a, b, c, d).
        """
        b = b if b is not None else make_zero_like(a)
        c = c if c is not None else make_zero_like(a)
        d = d if d is not None else make_zero_like(a)
        return a, b, c, d