import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)

import comfy.model_management as mm

from ..sampling.sampler import run_sampler
from ..utils.sampling_utils import prepare_conds


class MochiWrapperSamplerCustomNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MOCHIMODEL",),
                "positive": ("CONDITIONING", ),
                "negative": ("CONDITIONING", ),
                "cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 30.0, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "sigmas": ("SIGMAS", {"tooltip": "Override sigma schedule and steps"}),
                "latents": ("LATENT", ),
                "sampler": ("SAMPLER", ),
                "add_noise": ("BOOLEAN", ),
            }
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("samples",)
    FUNCTION = "process"
    CATEGORY = "MochiEdit/Wrapper"

    def process(self, model, positive, negative, cfg, seed, sigmas, latents, sampler, add_noise):
        mm.soft_empty_cache()

        sigmas = sigmas.tolist()
        latents = latents['samples']
        positive, negative = prepare_conds(positive, negative)
            
        latents = run_sampler(model, latents, positive, negative, sigmas, cfg, sampler.sampler_function, add_noise, seed)
    
        mm.soft_empty_cache()

        return ({"samples": latents},)