import torch
import gc


class Iter_ImageSwitch:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "output_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "input_count": ("INT", {"default": 2, "min": 1, "max": 64, "step": 1}),
                "use_last_input": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "switch_image"
    CATEGORY = "CustomNodes/Iterators"

    def switch_image(self, output_index, input_count, use_last_input, **kwargs):
        # 1. 取得目標索引（對 input_count 取模，確保不溢出）
        # 這樣 index 0~n 就會嚴格對應到 image_0 ~ image_n
        target_idx = output_index % input_count
        
        # 2. 從目標索引開始「往回」搜尋有效接點
        res_image = None
        for i in range(target_idx, -1, -1):
            val = kwargs.get(f"image_{i}")
            if val is not None or not use_last_input:
                res_image = val
                break
        
        gc.collect()
        torch.cuda.empty_cache() 
        return (res_image,)