import os
import torch
import folder_paths
from nodes import KSamplerAdvanced
import comfy.samplers
from safetensors.torch import save_file

from comfy_api.latest import io
try:
    from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
except ImportError:
    import nodes
    SamplerCustomAdvanced = nodes.NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]

class EvalSampler(KSamplerAdvanced):
    @classmethod
    def INPUT_TYPES(s):
        input_types = KSamplerAdvanced.INPUT_TYPES()
        new_inputs = {
            "required": dict(input_types["required"])
        }
        if "optional" in input_types:
            new_inputs["optional"] = dict(input_types["optional"])
            
        new_inputs["required"]["output_folder"] = ("STRING", {"default": "eval_latents"})
        return new_inputs

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "sample_eval"
    CATEGORY = "sampling"

    def sample_eval(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, output_folder, denoise=1.0):
        model = model.clone()
        
        out_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
        os.makedirs(out_dir, exist_ok=True)
        
        step_counter = [0]
        
        old_cfg_func = model.model_options.get("sampler_cfg_function", None)
        
        def custom_cfg_hook(args):
            cond_denoised = args["cond_denoised"]
            uncond_denoised = args["uncond_denoised"]
            
            step = step_counter[0]
            
            batch_size = cond_denoised.shape[0]
            for b_idx in range(batch_size):
                b_seed = noise_seed + b_idx
                
                cond_b = cond_denoised[b_idx:b_idx+1].contiguous()
                save_file({"latent": cond_b}, os.path.join(out_dir, f"{b_seed}_S{step}.safetensors"))
                
                uncond_b = uncond_denoised[b_idx:b_idx+1].contiguous()
                save_file({"latent": uncond_b}, os.path.join(out_dir, f"{b_seed}_S{step}_uncond.safetensors"))
            
            step_counter[0] += 1
            
            if old_cfg_func is not None:
                return old_cfg_func(args)
            else:
                cfg_result = args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]
                return args["input"] - cfg_result
        
        model.set_model_sampler_cfg_function(custom_cfg_hook)
        
        return super().sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=denoise)


class EvalSamplerAdvanced(SamplerCustomAdvanced):
    @classmethod
    def define_schema(cls):
        schema = super().define_schema()
        schema.node_id = "EvalSamplerAdvanced"
        schema.inputs.append(io.String.Input("output_folder", default="eval_latents"))
        return schema

    @classmethod
    def execute(cls, noise, guider, sampler, sigmas, latent_image, output_folder) -> io.NodeOutput:
        # Clone guider's model_patcher to avoid side effects on the original model
        guider.model_patcher = guider.model_patcher.clone()
        if hasattr(guider, "model_options"):
            guider.model_options = guider.model_patcher.model_options
        
        out_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
        os.makedirs(out_dir, exist_ok=True)
        
        step_counter = [0]
        noise_seed = getattr(noise, "seed", 0)
        
        old_cfg_func = guider.model_patcher.model_options.get("sampler_cfg_function", None)
        
        def custom_cfg_hook(args):
            cond_denoised = args["cond_denoised"]
            uncond_denoised = args["uncond_denoised"]
            
            step = step_counter[0]
            
            batch_size = cond_denoised.shape[0]
            for b_idx in range(batch_size):
                b_seed = noise_seed + b_idx
                
                cond_b = cond_denoised[b_idx:b_idx+1].contiguous()
                save_file({"latent": cond_b}, os.path.join(out_dir, f"{b_seed}_S{step}.safetensors"))
                
                uncond_b = uncond_denoised[b_idx:b_idx+1].contiguous()
                save_file({"latent": uncond_b}, os.path.join(out_dir, f"{b_seed}_S{step}_uncond.safetensors"))
            
            step_counter[0] += 1
            
            if old_cfg_func is not None:
                return old_cfg_func(args)
            else:
                cfg_result = args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]
                return args["input"] - cfg_result
        
        guider.model_patcher.set_model_sampler_cfg_function(custom_cfg_hook, disable_cfg1_optimization=True)
        
        return super().execute(noise, guider, sampler, sigmas, latent_image)