import comfy.samplers
import comfy.sample

from .hd_preview import create_preview_callback

def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative,
latent, preview_decode, vae, preview_mode,denoise = 1.0):
    latent_image = latent["samples"]
    latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
    noise = comfy.sample.prepare_noise(latent_image, seed, None)

    if preview_mode == "Sync":
        callback = create_preview_callback(model, vae, steps, preview_decode, interval = 1, sync_preview = True)
    else:
        callback = create_preview_callback(model, vae, steps, preview_decode, interval = 5, sync_preview = False)
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler,
    positive, negative, latent_image, denoise = denoise, callback = callback, seed = seed)
    out = latent.copy()
    out["samples"] = samples
    return (out, )

class HighPreviewKSampler:

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "model": ("MODEL", ),
            "seed": ("INT", {"default" : 0}),
            "steps": ("INT", {"default": 20}),
            "cfg": ("FLOAT", {"default": 8}),
            "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
            "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
            "positive": ("CONDITIONING", ),
            "negative": ("CONDITIONING", ),
            "latent_image": ("LATENT", ),
            "vae": ("VAE", ),
            "preview_mode": (["Sync", "Fast"], {"default": "Sync"}),
            "preview_decode": (["Full", "Scale"], {"default": "Scale"}),
            "denoise": ("FLOAT", {"default": 1.0, "min": 0, "max":1})
            }
        }
    RETURN_TYPES = ("LATENT", )
    FUNCTION = "sample"
    CATEGORY = "High_Preview"

    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative,
    latent_image, preview_decode, vae, preview_mode, denoise):
        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
            negative, latent_image, preview_decode, vae, preview_mode, denoise,)

NODE_CLASS_MAPPINGS = {"HighPreviewKSampler": HighPreviewKSampler}
NODE_DISPLAY_NAME_MAPPINGS = {"HighPreviewKSampler": "High Preview KSampler"}