import torch

from .zeval import LatentMap, parse


class ZEvalNode:
    """
    ComfyUI Custom Node with dynamic latent input count.
    """

    CATEGORY = "latent"
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "evaluate"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "expression": ("STRING", {"multiline": True, "default": "l1"}),
                "latent_1": ("LATENT",),
            },
        }

    @classmethod
    def IS_CHANGED(cls, expression, **kwargs):
        """Force re-evaluation when expression changes."""
        return expression or kwargs

    def evaluate(self, expression: str, **kwargs) -> tuple[dict[str, torch.Tensor]]:
        """
        Evaluate the expression using the provided latents.
        """

        latents = [
            (name[7:], latent['samples'])
            for name, latent in kwargs.items()
            if name.startswith("latent_")
        ]
        latents.sort()
        latent_map = LatentMap.create(latent for _, latent in latents)
        try:
            ast = parse(expression)
            print(f"ZEvalNode: evaluating expression {ast!r}")
            result = ast(latent_map)

        except Exception as e:
            raise RuntimeError(f"Latent expression parsing failed: {e}") from e

        return ({"samples": result},)