import random

class CausVidControl:
    CATEGORY = "utils"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "width": ("INT", {"default": 480, "min": 64, "max": 4096, "step": 16, "display": "number"}),
                "height": ("INT", {"default": 832, "min": 64, "max": 4096, "step": 16, "display": "number"}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "frames": ("INT", {"default": 81, "min": 1, "max": 1024, "step": 4, "display": "number"}),
                "steps": ("INT", {"default": 15, "min": 1, "max": 999, "step": 1, "display": "number"}),
                "clip": ("CLIP", ),
                "positive_prompt": ("STRING", {"multiline": True, "default": ""}),
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
            }
        }

    RETURN_TYPES = ("INT", "INT", "INT", "INT", "INT", "CONDITIONING", "CONDITIONING")
    RETURN_NAMES = ("width", "height", "seed", "frames", "steps", "positive_conditioning", "negative_conditioning")
    FUNCTION = "get_params"
    OUTPUT_NODE = True

    def get_params(self, width, height, seed, frames, steps, clip, positive_prompt, negative_prompt):
        positive_tokens = clip.tokenize(positive_prompt)
        positive_cond, positive_pooled = clip.encode_from_tokens(positive_tokens, return_pooled=True)
        positive_conditioning = [[positive_cond, {"pooled_output": positive_pooled}]]
        negative_tokens = clip.tokenize(negative_prompt)
        negative_cond, negative_pooled = clip.encode_from_tokens(negative_tokens, return_pooled=True)
        negative_conditioning = [[negative_cond, {"pooled_output": negative_pooled}]]
        return (width, height, seed, frames, steps, positive_conditioning, negative_conditioning)

NODE_CLASS_MAPPINGS = {
    "CausVidControl": CausVidControl
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CausVidControl": "CausVidControl"
}