import torch
import nodes # Импортируем стандартные ноды для доступа к common_ksampler
import comfy.sample

class BatchSeedSampler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "latent_image": ("LATENT",),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "control_after_generate": (["fixed", "increment", "decrement", "randomize"],),
                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

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

    def sample_batch(self, model, positive, negative, latent_image, seed, control_after_generate, steps, cfg, sampler_name, scheduler, denoise):
        samples = latent_image["samples"]
        batch_size = samples.shape[0]
        
        print(f"\n[BatchSampler] Начинаю обработку батча. Размер: {batch_size}")
        
        output_samples = []

        for i in range(batch_size):
            # Рассчитываем сид для текущего элемента
            current_seed = seed
            if control_after_generate == "increment":
                current_seed = seed + i
            elif control_after_generate == "decrement":
                current_seed = seed - i
            elif control_after_generate == "randomize":
                # Генерируем псевдорандом на основе базового сида и индекса
                current_seed = hash(f"{seed}_{i}") % 0xffffffffffffffff

            print(f"[BatchSampler] Обработка элемента {i+1}/{batch_size} с Seed: {current_seed}")

            # Выделяем i-й латент из батча (сохраняя размерность [1, C, H, W])
            single_latent = {"samples": samples[i:i+1]}

            # Используем системную функцию сэмплирования ComfyUI
            # Она возвращает кортеж, где первый элемент - это словарь с латентом
            try:
                result = nodes.common_ksampler(
                    model, 
                    current_seed, 
                    steps, 
                    cfg, 
                    sampler_name, 
                    scheduler, 
                    positive, 
                    negative, 
                    single_latent, 
                    denoise=denoise
                )
                output_samples.append(result[0]["samples"])
            except Exception as e:
                print(f"[BatchSampler] ОШИБКА на элементе {i}: {str(e)}")
                raise e

        # Объединяем все результаты обратно в один тензор батча
        final_samples = torch.cat(output_samples, dim=0)
        
        print(f"[BatchSampler] Завершено. Итоговый тензор: {final_samples.shape}\n")
        
        return ({"samples": final_samples},)