from .utils import convert_base64_to_image, convert_image_to_base64, convert_base64_to_audio

class Base64ToImage:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "base64": ("STRING", {"multiline": True})
            }
        }
    
    RETURN_TYPES = ("IMAGE","MASK",)
    RETURN_NAMES = ("image","mask",)

    CATEGORY = "Base64-Nodes"

    FUNCTION = "base64_to_image"

    def base64_to_image(self, base64):
        image, mask = convert_base64_to_image(base64)
        return (image, mask,)
    
class ImageToBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",)
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("base64",)

    CATEGORY = "Base64-Nodes"

    FUNCTION = "image_to_base64"

    def image_to_base64(self, image):
        base64 = convert_image_to_base64(image)
        return (base64,)
    
class AudioToBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "base64": ("STRING", {"multiline": True})
            },            
            "optional" : {
                "start_time": ("FLOAT", {"default": 0, "min": 0,}),
                "duration": ("FLOAT" , {"default": 0, "min": 0, "max": 10000000, "step": 0.01}),
            }
        }
    
    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("audio",)

    CATEGORY = "Base64-Nodes"

    FUNCTION = "base64_to_audio"
    
    def base64_to_audio(self, base64, start_time=0, duration=0):
        audio = convert_base64_to_audio(base64, start_time, duration)
        return (audio,)

class LastFrameToBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",)
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("base64",)

    CATEGORY = "Base64-Nodes"

    FUNCTION = "last_frame_to_base64"

    def last_frame_to_base64(self, images):
        base64 = convert_image_to_base64(images[-1])
        return (base64,)
    
class FirstFrameToBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",)
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("base64",)

    CATEGORY = "Base64-Nodes"

    FUNCTION = "first_frame_to_base64"

    def first_frame_to_base64(self, images):
        base64 = convert_image_to_base64(images[0])
        return (base64,)
    
class SaveBase64Websocket:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base64": ("STRING",)
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "save_base64"
    OUTPUT_NODE = True
    CATEGORY = "Base64-Nodes"

    def save_base64(self, base64):
        return {
            "ui": {
                "base64": [base64]
            }
        }