import os
import glob
import json
import random
import time
from aiohttp import web
from PIL import Image, ImageOps, ImageSequence  # type: ignore
import numpy as np  # type: ignore
import torch  # type: ignore
import folder_paths  # type: ignore
from server import PromptServer  # type: ignore

SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tiff", ".avif", ".tga"}

def get_input_directory():
    if folder_paths and hasattr(folder_paths, "get_input_directory"):
        return folder_paths.get_input_directory()
    return os.path.abspath("input")

def resolve_image_path(filename, subfolder=""):
    input_dir = get_input_directory()
    if subfolder:
        clean_sub = subfolder.strip("/\\")
        candidate = os.path.join(input_dir, clean_sub, filename)
    else:
        candidate = os.path.join(input_dir, filename)

    if os.path.exists(candidate) and os.path.isfile(candidate):
        return candidate

    if folder_paths and hasattr(folder_paths, "get_annotated_filepath"):
        try:
            annotated = folder_paths.get_annotated_filepath(filename)
            if os.path.exists(annotated):
                return annotated
        except Exception:
            pass

    # Direct fallback if filename is absolute or relative
    if os.path.exists(filename) and os.path.isfile(filename):
        return os.path.abspath(filename)

    return None

def scan_input_images(subfolder=""):
    input_dir = get_input_directory()
    target_dir = os.path.join(input_dir, subfolder.strip("/\\")) if subfolder else input_dir

    if not os.path.exists(target_dir) or not os.path.isdir(target_dir):
        return []

    results = []
    try:
        for root, _, files in os.walk(target_dir):
            for file in files:
                ext = os.path.splitext(file)[1].lower()
                if ext in SUPPORTED_EXTENSIONS:
                    full_path = os.path.join(root, file)
                    rel_to_input = os.path.relpath(full_path, input_dir).replace("\\", "/")
                    rel_dir = os.path.dirname(rel_to_input).replace("\\", "/")
                    file_size = os.path.getsize(full_path)
                    
                    results.append({
                        "filename": file,
                        "subfolder": rel_dir if rel_dir != "." else "",
                        "relative_path": rel_to_input,
                        "size": file_size,
                        "mtime": os.path.getmtime(full_path)
                    })
    except Exception as e:
        print(f"[EasyMultiImageLoader] Error scanning folder {target_dir}: {e}")

    results.sort(key=lambda x: x["relative_path"].lower())
    return results

def create_placeholder_image(text="No Image Loaded"):
    # Create a 512x512 dark grey tensor
    img_array = np.full((512, 512, 3), 0.15, dtype=np.float32)
    # Add a border
    img_array[0:8, :, :] = 0.3
    img_array[-8:, :, :] = 0.3
    img_array[:, 0:8, :] = 0.3
    img_array[:, -8:, :] = 0.3
    # Add subtle cross lines
    for i in range(512):
        if 0 <= i < 512:
            img_array[i, i, :] = 0.25
            img_array[i, 511 - i, :] = 0.25

    tensor = torch.from_numpy(img_array)[None,]
    mask = torch.zeros((1, 512, 512), dtype=torch.float32)
    return tensor, mask

def load_single_image(filepath):
    try:
        img = Image.open(filepath)
        output_images = []
        output_masks = []
        
        for frame in ImageSequence.Iterator(img):
            frame = ImageOps.exif_transpose(frame)
            if frame.mode == 'I':
                frame = frame.point(lambda i: i * (1 / 255))
            
            # Mask extraction
            if 'A' in frame.getbands():
                mask = np.array(frame.getchannel('A')).astype(np.float32) / 255.0
                mask = 1.0 - torch.from_numpy(mask)
            else:
                mask = torch.zeros((frame.height, frame.width), dtype=torch.float32)

            rgb_frame = frame.convert("RGB")
            np_frame = np.array(rgb_frame).astype(np.float32) / 255.0
            tensor_frame = torch.from_numpy(np_frame)[None,]

            output_images.append(tensor_frame)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1:
            out_img = torch.cat(output_images, dim=0)
            out_mask = torch.cat(output_masks, dim=0)
        else:
            out_img = output_images[0]
            out_mask = output_masks[0]

        return out_img, out_mask
    except Exception as e:
        print(f"[EasyMultiImageLoader] Error loading image {filepath}: {e}")
        return None, None

# ==================== API ENDPOINTS ====================

if PromptServer is not None:
    @PromptServer.instance.routes.get("/easy_multi_image/get_input_files")
    async def api_get_input_files(request):
        try:
            subfolder = request.query.get("subfolder", "")
            files = scan_input_images(subfolder)
            return web.json_response({"files": files, "input_directory": get_input_directory()})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    @PromptServer.instance.routes.post("/easy_multi_image/scan_subfolder")
    async def api_scan_subfolder(request):
        try:
            data = await request.json()
            subfolder = data.get("subfolder", "")
            files = scan_input_images(subfolder)
            return web.json_response({"files": files})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)


# ==================== COMFYUI CUSTOM NODE ====================

def extract_node_id(unique_id):
    if isinstance(unique_id, (list, tuple)):
        return str(unique_id[0]) if len(unique_id) > 0 else "default_node"
    if unique_id is not None:
        return str(unique_id)
    return "default_node"

class EasyMultiImageLoader:
    """
    ComfyUI node to load multiple images from ComfyUI's standard input directory,
    manually choose which image to output, or automatically cycle through images with
    smooth wrap-around and progress reporting.
    """
    
    # Internal in-memory tracking of cycle steps per node instance
    _node_cycle_state = {}

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (["Fixed / Selected Index", "Cycle (Increment)", "Cycle (Decrement)", "Random"], {"default": "Fixed / Selected Index"}),
                "index": ("INT", {"default": 1, "min": 1, "max": 999999, "step": 1}),
            },
            "optional": {
                "images_list": ("STRING", {"default": "[]", "multiline": True}),
                "subfolder": ("STRING", {"default": "", "multiline": False}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "extra_pnginfo": "EXTRA_PNGINFO",
                "prompt": "PROMPT",
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "INT", "INT", "FLOAT")
    RETURN_NAMES = ("IMAGE", "MASK", "filename", "current_index", "total_images", "progress_percent")
    FUNCTION = "process_images"
    CATEGORY = "🗄️ Comfy Cabinet"

    @classmethod
    def IS_CHANGED(cls, mode, index, images_list="[]", subfolder="", batch_size=1, unique_id=None, **kwargs):
        # If cycling or random, always return unique timestamp so ComfyUI executes on every run in the queue
        if mode in ("Cycle (Increment)", "Cycle (Decrement)", "Random"):
            return time.time_ns()
        # Otherwise hash the static inputs
        return f"{mode}_{index}_{images_list}_{subfolder}_{batch_size}"

    def process_images(self, mode, index, images_list="[]", subfolder="", batch_size=1, unique_id=None, extra_pnginfo=None, prompt=None):
        # Parse images list
        parsed_items = []
        if images_list and images_list.strip():
            try:
                raw_json = json.loads(images_list)
                if isinstance(raw_json, list):
                    for item in raw_json:
                        if isinstance(item, str) and item.strip():
                            parsed_items.append(item.strip())
                        elif isinstance(item, dict) and "relative_path" in item:
                            parsed_items.append(item["relative_path"])
                        elif isinstance(item, dict) and "filename" in item:
                            sf = item.get("subfolder", "")
                            fn = item["filename"]
                            parsed_items.append(os.path.join(sf, fn).replace("\\", "/") if sf else fn)
            except Exception:
                # Fallback: line-delimited or comma-delimited strings
                for line in images_list.replace("\r", "").split("\n"):
                    for chunk in line.split(","):
                        c = chunk.strip().strip('"').strip("'")
                        if c:
                            parsed_items.append(c)

        # If images_list was empty but a subfolder is specified, auto-scan that subfolder
        if not parsed_items and subfolder:
            scanned = scan_input_images(subfolder)
            parsed_items = [s["relative_path"] for s in scanned]

        total_images = len(parsed_items)

        if total_images == 0:
            print("[EasyMultiImageLoader] Warning: No images found or loaded in the list.")
            placeholder_img, placeholder_mask = create_placeholder_image("No Images Loaded")
            return {
                "ui": {
                    "index": [0],
                    "filename": [""],
                    "progress_percent": [0.0]
                },
                "result": (placeholder_img, placeholder_mask, "", 0, 0, 0.0)
            }

        # Determine effective target index (1-based)
        node_key = extract_node_id(unique_id)
        state = self._node_cycle_state.get(node_key, {})
        prev_idx = state.get("index", None)
        prev_mode = state.get("mode", None)

        if mode == "Fixed / Selected Index":
            target_index = ((index - 1) % total_images) + 1
            self._node_cycle_state[node_key] = {"index": target_index, "mode": mode}
        elif mode == "Cycle (Increment)":
            if prev_idx is None or prev_mode != mode:
                target_index = ((index - 1) % total_images) + 1
            else:
                target_index = (prev_idx % total_images) + 1
            self._node_cycle_state[node_key] = {"index": target_index, "mode": mode}
        elif mode == "Cycle (Decrement)":
            if prev_idx is None or prev_mode != mode:
                target_index = ((index - 1) % total_images) + 1
            else:
                target_index = ((prev_idx - 2) % total_images) + 1
            self._node_cycle_state[node_key] = {"index": target_index, "mode": mode}
        elif mode == "Random":
            target_index = random.randint(1, total_images)
            self._node_cycle_state[node_key] = {"index": target_index, "mode": mode}
        else:
            target_index = max(1, min(index, total_images))
            self._node_cycle_state[node_key] = {"index": target_index, "mode": mode}

        # Calculate progress percentage (100% when at last image)
        progress_percent = round((target_index / total_images) * 100.0, 2)

        # Collect images for batch
        batch_images = []
        batch_masks = []
        filenames_loaded = []

        actual_batch = max(1, min(batch_size, total_images))
        ref_h, ref_w = None, None

        for b in range(actual_batch):
            # Wrap around batch indices if batch_size > 1
            curr_b_idx = ((target_index - 1 + b) % total_images) + 1
            img_rel_path = parsed_items[curr_b_idx - 1]
            full_path = resolve_image_path(img_rel_path, subfolder)

            if not full_path:
                print(f"[EasyMultiImageLoader] Warning: Could not resolve file path for '{img_rel_path}'")
                continue

            img_tensor, mask_tensor = load_single_image(full_path)
            if img_tensor is None:
                continue

            # Ensure mask is a valid 3D tensor [B, H, W]
            if mask_tensor is None:
                mask_tensor = torch.zeros((img_tensor.shape[0], img_tensor.shape[1], img_tensor.shape[2]), dtype=torch.float32)
            elif mask_tensor.ndim == 2:
                mask_tensor = mask_tensor.unsqueeze(0)

            # Ensure all batch items match first image's spatial dimensions (H, W)
            if ref_h is None:
                ref_h, ref_w = img_tensor.shape[1], img_tensor.shape[2]
            else:
                if img_tensor.shape[1] != ref_h or img_tensor.shape[2] != ref_w:
                    # Reshape / interpolate to match reference
                    img_perm = img_tensor.permute(0, 3, 1, 2)  # [B, C, H, W]
                    img_resized = torch.nn.functional.interpolate(img_perm, size=(ref_h, ref_w), mode="bilinear", align_corners=False)
                    img_tensor = img_resized.permute(0, 2, 3, 1)

                    mask_unsqueeze = mask_tensor.unsqueeze(1)  # [B, 1, H, W]
                    mask_resized = torch.nn.functional.interpolate(mask_unsqueeze, size=(ref_h, ref_w), mode="bilinear", align_corners=False)
                    mask_tensor = mask_resized.squeeze(1)

            batch_images.append(img_tensor)
            batch_masks.append(mask_tensor)
            filenames_loaded.append(os.path.basename(img_rel_path))

        if not batch_images:
            placeholder_img, placeholder_mask = create_placeholder_image("Failed to Load Image")
            return {
                "ui": {
                    "index": [target_index],
                    "filename": [""],
                    "progress_percent": [progress_percent]
                },
                "result": (placeholder_img, placeholder_mask, "", target_index, total_images, progress_percent)
            }

        final_image = torch.cat(batch_images, dim=0)
        final_mask = torch.cat(batch_masks, dim=0)
        primary_filename = parsed_items[target_index - 1]

        # Send live progress update event to ComfyUI canvas node
        if PromptServer is not None:
            try:
                PromptServer.instance.send_sync("easy_multi_image_progress", {
                    "node_id": node_key,
                    "current_index": target_index,
                    "total_images": total_images,
                    "progress_percent": progress_percent,
                    "filename": primary_filename,
                    "mode": mode
                })
            except Exception as e:
                pass

        return {
            "ui": {
                "index": [target_index],
                "filename": [primary_filename],
                "progress_percent": [progress_percent]
            },
            "result": (final_image, final_mask, primary_filename, target_index, total_images, progress_percent)
        }


NODE_CLASS_MAPPINGS = {
    "EasyMultiImageLoader": EasyMultiImageLoader
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "EasyMultiImageLoader": "🖼️ Easy Multi-Image Loader"
}