import os
import time
import torch
import folder_paths
from server import PromptServer
from aiohttp import web
from PIL import Image
import numpy as np

CACHE_DIR = os.path.join(folder_paths.get_temp_directory(), "image_chooser_cache")
os.makedirs(CACHE_DIR, exist_ok=True)

# Глобальный словарь для связи бэкенда и фронтенда
CHOICES = {}

class ImageChooser:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_A": ("IMAGE",),
                "image_B": ("IMAGE",),
            },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "my_id": "UNIQUE_ID"},
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "execute"
    CATEGORY = "utils"

    # ЭТОТ МЕТОД РЕШАЕТ ПРОБЛЕМУ КЭША: 
    # Возвращает текущее время. ComfyUI видит, что значение изменилось, 
    # аннулирует кэш этой ноды и заставляет её запускаться при каждом Queue Prompt.
    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return time.time()

    def execute(self, image_A, image_B, prompt=None, extra_pnginfo=None, my_id=None):
        my_id_str = str(my_id)
        
        # Перезаписываем картинки свежими кадрами
        self.save_temp_image(image_A, f"{my_id_str}_A.png")
        self.save_temp_image(image_B, f"{my_id_str}_B.png")

        # Обязательно сбрасываем старый выбор пользователя для этой конкретной ноды
        CHOICES[my_id_str] = None

        # Отправляем сигнал браузеру открыть модальное окно
        PromptServer.instance.send_sync("image_chooser_wait", {"node_id": my_id_str})

        print(f"[Image Chooser] Workflow paused. Waiting for user selection on node {my_id_str}...")
        
        # Ждем клика
        while CHOICES.get(my_id_str) is None:
            time.sleep(0.1)

        choice = CHOICES.pop(my_id_str)
        
        print(f"[Image Chooser] Pause lifted. Result received: {choice}")

        if choice == "CANCEL":
            raise Exception("Execution cancelled by user via Image Chooser.")

        if choice == "A":
            return (image_A,)
        else:
            return (image_B,)

    def save_temp_image(self, tensor, filename):
        if len(tensor.shape) == 4:
            img_np = tensor[0].cpu().numpy()
        else:
            img_np = tensor.cpu().numpy()
            
        img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
        img = Image.fromarray(img_np)
        img.save(os.path.join(CACHE_DIR, filename))


@PromptServer.instance.routes.post("/image_chooser/select")
async def handle_select(request):
    json_data = await request.json()
    node_id = str(json_data.get("node_id"))
    choice = json_data.get("choice")
    
    CHOICES[node_id] = choice
    return web.json_response({"status": "success"})


NODE_CLASS_MAPPINGS = {
    "ImageChooser": ImageChooser
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ImageChooser": "Image Chooser"
}

WEB_DIRECTORY = "."