import sys
import os
sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))

import nodes
import folder_paths


class LoadOffset:
    @classmethod
    def INPUT_TYPES(s):
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        return {
            "required": {
                "image": (sorted(files), {"image_upload": True}),
                "vae": ("VAE",),
                "width":("INT", {
                    "default": 512,
                    "min": 256,
                    "max": 2048,
                    "step": 64
                }),
                "height": ("INT", {
                    "default": 512,
                    "min": 256,
                    "max": 2048,
                    "step": 64
                }),
                "repeat": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 128,
                    "step": 1
                })
            },
        }

    CATEGORY = "latent"

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "load_image"
    def load_image(self, image, vae, width, height, repeat):
        loader = nodes.LoadImage()
        scaler = nodes.ImageScale()
        encoder = nodes.VAEEncode()
        repeater = nodes.RepeatLatentBatch()

        img,_ = loader.load_image(image)
        img, = scaler.upscale(img, "bicubic", width, height, False)
        latent, = encoder.encode(vae, img)
        latents, = repeater.repeat(latent, repeat)

        return latents,

NODE_CLASS_MAPPINGS = {
    "LoadOffset": LoadOffset,
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadOffset": "Load image as latent",
}