"""GridImagePool — the Image Pool (Grid) ComfyUI node."""
import os
from .gates_compat import grid_pool_base as _grid_pool_base
from . import pool, imaging

NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}


class GridImagePool:
    CATEGORY = "Datasete Gates"
    FUNCTION = "run"
    RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
    RETURN_NAMES = ("image", "mask", "index", "count", "label")

    @classmethod
    def INPUT_TYPES(cls):
        # pool_id is a per-node UUID owned by the JS extension. It must be a
        # NORMAL input, not a "hidden" one: ComfyUI only populates hidden inputs
        # for built-in types (UNIQUE_ID, PROMPT, ...), so a custom hidden
        # "POOL_ID" would always arrive as None. The frontend, however, forwards
        # every serialized widget value by name, so a declared STRING input
        # backed by a (hidden-rendered) widget reliably reaches run().
        return {
            "required": {
                "index": ("INT", {"default": -1, "min": -1, "max": 9999}),
                "pool_id": ("STRING", {"default": "default"}),
            },
            # optional companion input: a Pool Profile node feeds the profile id
            # here; when connected it overrides pool_id (see `effective` below).
            "optional": {
                "profile": ("POOL_PROFILE",),
            },
        }

    @staticmethod
    def _resolve(index, pool_id):
        base = _grid_pool_base()
        m = pool.read_manifest(base, pool_id)
        idx = pool.resolve_slot(m, index)
        return base, m, idx

    def run(self, index, pool_id="default", profile=None):
        effective = profile or pool_id
        base, m, idx = self._resolve(index, effective)
        if idx < 0:
            img, mask = imaging.empty_outputs()
            return (img, mask, 0, 0, "")
        slot = m["slots"][idx]
        d = pool.pool_dir(base, effective)
        img = imaging.load_image_tensor(str(d / slot["image"]))
        h, w = int(img.shape[1]), int(img.shape[2])
        mask_name = slot.get("mask")
        mask = imaging.load_mask_tensor(str(d / mask_name) if mask_name else None, h, w)
        return (img, mask, idx, len(m["slots"]), slot.get("label", ""))

    @classmethod
    def IS_CHANGED(cls, index, pool_id="default", profile=None, **kwargs):
        effective = profile or pool_id
        base, m, idx = cls._resolve(index, effective)
        if idx < 0:
            return imaging.change_hash(effective, -1, [])
        slot = m["slots"][idx]
        d = pool.pool_dir(base, effective)
        mtimes = []
        for key in ("image", "mask"):
            name = slot.get(key)
            p = d / name if name else None
            mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0)
        # include active so manual selection changes invalidate cache
        return imaging.change_hash(effective, f"{idx}:{m.get('active')}", mtimes)


NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool}
NODE_DISPLAY_NAME_MAPPINGS = {"GridImagePool": "Image Pool (Grid)"}