import json
import os
import re
import time
import uuid
from typing import Any, Dict, Optional
import numpy as np
import folder_paths  # pyright: ignore[reportMissingImports]
from ..services.service_registry import ServiceRegistry
from ..metadata_collector.metadata_processor import MetadataProcessor
from ..metadata_collector import get_metadata
from ..utils.constants import CARD_PREVIEW_WIDTH
from ..utils.exif_utils import ExifUtils
from ..utils.utils import calculate_recipe_fingerprint, sanitize_folder_name
from PIL import Image, PngImagePlugin
import piexif  # pyright: ignore[reportMissingTypeStubs]
import logging

# Civitai-compatible sampler name mapping: ComfyUI internal → A1111 display name
CIVITAI_SAMPLER_MAP = {
    "euler": "Euler",
    "euler_ancestral": "Euler a",
    "lms": "LMS",
    "heun": "Heun",
    "dpm_2": "DPM2",
    "dpm_2_ancestral": "DPM2 a",
    "dpmpp_2s_ancestral": "DPM++ 2S a",
    "dpmpp_2m": "DPM++ 2M",
    "dpmpp_sde": "DPM++ SDE",
    "dpmpp_sde_gpu": "DPM++ SDE",
    "dpmpp_2m_sde": "DPM++ 2M SDE",
    "dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
    "dpmpp_3m_sde": "DPM++ 3M SDE",
    "dpm_fast": "DPM fast",
    "dpm_adaptive": "DPM adaptive",
    "ddim": "DDIM",
    "plms": "PLMS",
    "uni_pc_bh2": "UniPC",
    "uni_pc": "UniPC",
    "lcm": "LCM",
}

# Base model display name → AIR URN slug
# Sourced from civitai source: src/shared/constants/basemodel.constants.ts
BASE_MODEL_AIR_SLUG = {
    # Stable Diffusion family
    "SD 1.4": "sd1",
    "SD 1.5": "sd1",
    "SD 1.5 LCM": "sd1",
    "SD 1.5 Hyper": "sd1",
    "SD 2.0": "sd2",
    "SD 2.0 768": "sd2",
    "SD 2.1": "sd2",
    "SD 2.1 768": "sd2",
    "SD 2.1 Unclip": "sd2",
    "SD 3.0": "sd3",
    "SD 3.5": "sd35",
    "SD 3.5 Large": "sd35",
    "SD 3.5 Large Turbo": "sd35",
    "SD 3.5 Medium": "sd35",
    "SDXL 0.9": "sdxl",
    "SDXL 1.0": "sdxl",
    "SDXL 1.0 LCM": "sdxl",
    "SDXL Lightning": "sdxl",
    "SDXL Hyper": "sdxl",
    "SDXL Turbo": "sdxl",
    "SDXL Distilled": "sdxldistilled",
    "Stable Cascade": "scascade",
    "Stable Video Diffusion": "svd",
    "SVD": "svd",
    "SVD XT": "svdxt",

    # SDXL community fine-tunes
    "Pony": "pony",
    "Pony Diffusion": "pony",
    "Illustrious": "illustrious",
    "NoobAI": "noobai",
    "Animagine": "illustrious",

    # Flux family
    "Flux.1": "flux1",
    "Flux.1 D": "flux1",
    "Flux.1 S": "flux1",
    "Flux.1 Krea": "fluxkrea",
    "Flux.1 Kontext": "flux1kontext",
    "Flux.2": "flux2",
    "Flux.2 D": "flux2",
    "Flux.2 Klein 9B": "flux2klein_9b",
    "Flux.2 Klein 9B Base": "flux2klein_9b_base",
    "Flux.2 Klein 4B": "flux2klein_4b",
    "Flux.2 Klein 4B Base": "flux2klein_4b_base",

    # Other image models (sorted alphabetically)
    "AuraFlow": "auraflow",
    "Chroma": "chroma",
    "HiDream": "hidream",
    "HiDream-O1": "hidream-o1",
    "Hunyuan DiT": "hydit1",
    "Hunyuan Video": "hyv1",
    "Kolors": "kolors",
    "Lumina": "lumina",
    "Mochi": "mochi",
    "ODOR": "odor",
    "PixArt Alpha": "pixarta",
    "PixArt Sigma": "pixarte",
    "Playground v2": "playgroundv2",
    "Playground v2.5": "playgroundv2",
    "Pony Diffusion V7": "ponyv7",

    # Video models
    "CogVideoX": "cogvideox",
    "LTX Video": "ltxv",
    "LTX Video 2": "ltxv2",
    "LTX Video 2.3": "ltxv23",
    "Wan Video": "wanvideo",
    "Wan Video 1.3B T2V": "wanvideo_13b_t2v",
    "Wan Video 14B T2V": "wanvideo_14b_t2v",
    "Wan Video 14B I2V 480p": "wanvideo_14b_i2v_480p",
    "Wan Video 14B I2V 720p": "wanvideo_14b_i2v_720p",

    # Third-party / proprietary image models
    "Boogu": "boogu",
    "Ernie": "ernie",
    "Grok": "grok",
    "HappyHorse": "happyhorse",
    "Ideogram": "ideogram",
    "Ideogram 4.0": "ideogram",
    "Imagen": "imagen4",
    "Imagen 4": "imagen4",
    "Krea": "krea2",
    "Krea 2": "krea2",
    "Lens": "lens",
    "MAI": "mai",
    "Nano Banana": "nanobanana",
    "OpenAI": "openai",
    "Reve": "reve",
    "Reve 2": "reve",
    "Reve 2.1": "reve",
    "Seedream": "seedream",
    "Sora": "sora2",
    "Sora 2": "sora2",
    "Veo": "veo3",
    "Veo 2": "veo3",
    "Veo 3": "veo3",
    "ZImageTurbo": "zimageturbo",
    "ZImageBase": "zimagebase",
    "ZImage": "zimagebase",

    # Third-party video models
    "Hailuo by MiniMax": "minimax",
    "Haiper": "haiper",
    "Kling": "kling",
    "Lightricks": "lightricks",
    "Seedance": "seedance",
    "Vidu": "vidu",

    # Qwen family
    "Qwen": "qwen",
    "Qwen 2": "qwen2",

    # Anima
    "Anima": "anima",

    # Special
    "Upscaler": "upscaler",
    "Other": "other",
}

logger = logging.getLogger(__name__)


class SaveImageLM:
    NAME = "Save Image (LoraManager)"
    CATEGORY = "Lora Manager/utils"
    DESCRIPTION = "Save images with embedded generation metadata in compatible format"

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.prefix_append = ""
        self.compress_level = 4
        self.counter = 0

    # Add pattern format regex for filename substitution
    pattern_format = re.compile(r"(%[^%]+%)")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "filename_prefix": (
                    "STRING",
                    {
                        "default": "ComfyUI",
                        "tooltip": "Base filename for saved images. Supports format patterns like %seed%, %width%, %height%, %model%, etc.",
                    },
                ),
                "file_format": (
                    ["png", "jpeg", "webp"],
                    {
                        "tooltip": "Image format to save as. PNG preserves quality, JPEG is smaller, WebP balances size and quality."
                    },
                ),
            },
            "optional": {
                "lossless_webp": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "When enabled, saves WebP images with lossless compression. Results in larger files but no quality loss.",
                    },
                ),
                "quality": (
                    "INT",
                    {
                        "default": 100,
                        "min": 1,
                        "max": 100,
                        "tooltip": "Compression quality for JPEG and lossy WebP formats (1-100). Higher values mean better quality but larger files.",
                    },
                ),
                "webp_method": (
                    "INT",
                    {
                        "default": 6,
                        "min": 0,
                        "max": 6,
                        "tooltip": "WebP compression method (0-6). 0=fastest/largest, 6=slowest/smallest. Only applies when file_format is 'webp'.",
                    },
                ),
                "jpeg_subsampling": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 2,
                        "tooltip": "JPEG chroma subsampling level. 0=4:4:4 (best quality), 1=4:2:2, 2=4:2:0 (smallest files). Only applies when file_format is 'jpeg'.",
                    },
                ),
                "embed_workflow": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "When enabled, saved images store the complete workflow. Drag the image back into ComfyUI to restore the original node graph. PNG and WebP only.",
                    },
                ),
                "save_with_metadata": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "When enabled, embeds generation parameters into the saved image metadata. Disable to skip writing generation metadata.",
                    },
                ),
                "add_loras_to_prompt": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "When enabled, appends the LoRA syntax line (e.g. <lora:name:strength>) after the positive prompt in the saved metadata.",
                    },
                ),
                "add_counter_to_filename": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Adds an incremental counter to filenames to prevent overwriting previous images.",
                    },
                ),
                "save_as_recipe": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Also saves each generated image as a LoRA Manager recipe.",
                    },
                ),
            },
            "hidden": {
                "id": "UNIQUE_ID",
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "process_image"
    OUTPUT_NODE = True

    def get_lora_hash(self, lora_name):
        """Get the lora hash from cache"""
        scanner = ServiceRegistry.get_service_sync("lora_scanner")

        # Use the new direct filename lookup method
        if scanner is not None:
            hash_value = scanner.get_hash_by_filename(lora_name)
            if hash_value:
                return hash_value

        return None

    def get_checkpoint_hash(self, checkpoint_path):
        """Get the checkpoint hash from cache"""
        scanner = ServiceRegistry.get_service_sync("checkpoint_scanner")

        if not checkpoint_path:
            return None

        # Extract basename without extension
        checkpoint_name = os.path.basename(checkpoint_path)
        checkpoint_name = os.path.splitext(checkpoint_name)[0]

        # Try direct filename lookup first
        if scanner is not None:
            hash_value = scanner.get_hash_by_filename(checkpoint_name)
            if hash_value:
                return hash_value

        return None

    def _resolve_model_cache_entry(self, scanner_type: str, name: str):
        """Resolve model hash, civitai metadata, and base_model from scanner cache.
        Returns (hash_str, civitai_dict, base_model_str). All values are empty defaults when not found."""
        scanner = ServiceRegistry.get_service_sync(scanner_type)
        if scanner is None or not name:
            return "", {}, ""

        entry = self._get_cached_model_by_name(scanner, name)
        if entry is None:
            basename = os.path.splitext(os.path.basename(name))[0]
            hash_val = scanner.get_hash_by_filename(basename)
            return (hash_val or "").lower(), {}, ""

        hash_val = (entry.get("sha256") or "").lower()
        civitai = entry.get("civitai") or {}
        base_model = entry.get("base_model") or ""
        return hash_val, civitai, base_model

    @staticmethod
    def _get_civitai_sampler_name(sampler_name: str, scheduler: str) -> str:
        if sampler_name in CIVITAI_SAMPLER_MAP:
            civitai_name = CIVITAI_SAMPLER_MAP[sampler_name]
            if scheduler == "karras":
                civitai_name += " Karras"
            elif scheduler == "exponential":
                civitai_name += " Exponential"
            return civitai_name
        else:
            if scheduler and scheduler != "normal":
                return f"{sampler_name}_{scheduler}"
            return sampler_name

    @staticmethod
    def _build_air_string(base_model: str, model_type: str, model_id: int, version_id: int) -> str:
        slug = BASE_MODEL_AIR_SLUG.get(base_model, "other")
        type_lower = model_type.lower() if model_type else "other"
        return f"urn:air:{slug}:{type_lower}:civitai:{model_id}@{version_id}"

    def format_metadata(self, metadata_dict: dict[str, Any], add_loras_to_prompt: bool = False) -> str:
        """Format metadata as A1111-compatible parameters string with Hashes JSON and Civitai resources."""
        if not metadata_dict: return ""

        prompt = metadata_dict.get("prompt", "")
        negative_prompt = metadata_dict.get("negative_prompt", "")
        steps = metadata_dict.get("steps")
        cfg = metadata_dict.get("guidance")
        if cfg is None:
            cfg = metadata_dict.get("cfg_scale")
        if cfg is None:
            cfg = metadata_dict.get("cfg")
        seed = metadata_dict.get("seed")
        size = metadata_dict.get("size")
        sampler = metadata_dict.get("sampler") or ""
        scheduler = metadata_dict.get("scheduler") or "normal"
        checkpoint = metadata_dict.get("checkpoint") or ""
        loras_text = metadata_dict.get("loras", "")
        clip_skip = metadata_dict.get("clip_skip")

        # Parse LoRA entries from <lora:name:strength> format
        lora_entries: list[tuple[str, float]] = []
        if loras_text:
            for match in re.findall(r"<lora:([^:]+):([^>]+)>", loras_text):
                lora_name, strength_str = match
                try:
                    strength = float(strength_str)
                except (ValueError, TypeError):
                    strength = 1.0
                lora_entries.append((lora_name, strength))

        # Resolve checkpoint hash and Civitai data from local cache
        ckpt_hash, ckpt_civitai, ckpt_base_model = "", {}, ""
        ckpt_display_name = ""
        if checkpoint:
            ckpt_hash, ckpt_civitai, ckpt_base_model = self._resolve_model_cache_entry(
                "checkpoint_scanner", checkpoint
            )
            ckpt_display_name = os.path.splitext(os.path.basename(checkpoint))[0]

        # Resolve LoRA hash and Civitai data from local cache
        loras_data: list[dict[str, Any]] = []
        for lora_name, strength in lora_entries:
            lora_hash, lora_civitai, lora_base_model = self._resolve_model_cache_entry(
                "lora_scanner", lora_name
            )
            loras_data.append({
                "name": lora_name,
                "strength": strength,
                "hash": lora_hash,
                "civitai": lora_civitai,
                "base_model": lora_base_model,
            })

        # Build Hashes JSON (A1111 / Civitai standard format)
        hashes: dict[str, str] = {}
        if ckpt_hash:
            hashes["model"] = ckpt_hash[:10].upper()
        for lora in loras_data:
            if lora["hash"]:
                hashes[f"LORA:{lora['name']}"] = lora["hash"][:10].upper()

        # Build Civitai resources JSON array
        civitai_resources: list[dict[str, Any]] = []
        if ckpt_civitai.get("id", 0) > 0:
            ckpt_resource: dict[str, Any] = {}
            ckpt_type = (ckpt_civitai.get("model") or {}).get("type", "Checkpoint")
            model_id = ckpt_civitai.get("modelId", 0)
            version_id = ckpt_civitai.get("id", 0)
            if model_id and version_id:
                ckpt_resource["air"] = self._build_air_string(
                    ckpt_base_model, ckpt_type, int(model_id), int(version_id)
                )
            elif version_id:
                ckpt_resource["modelVersionId"] = int(version_id)
            if ckpt_civitai.get("name"):
                ckpt_resource["versionName"] = ckpt_civitai["name"]
            if ckpt_resource:
                civitai_resources.append(ckpt_resource)

        for lora in loras_data:
            lora_civitai = lora["civitai"]
            if not lora_civitai or lora_civitai.get("id", 0) <= 0:
                continue
            lora_resource: dict[str, Any] = {"weight": lora["strength"]}
            lora_type = (lora_civitai.get("model") or {}).get("type", "LORA")
            model_id = lora_civitai.get("modelId", 0)
            version_id = lora_civitai.get("id", 0)
            if model_id and version_id:
                lora_resource["air"] = self._build_air_string(
                    lora["base_model"], lora_type, int(model_id), int(version_id)
                )
            elif version_id:
                lora_resource["modelVersionId"] = int(version_id)
            if lora_civitai.get("name"):
                lora_resource["versionName"] = lora_civitai["name"]
            civitai_resources.append(lora_resource)

        sampler_name = CIVITAI_SAMPLER_MAP.get(sampler, sampler) if sampler else None

        scheduler_mapping = {
            "normal": "Normal",
            "karras": "Karras",
            "exponential": "Exponential",
            "sgm_uniform": "SGM Uniform",
            "sgm_quadratic": "SGM Quadratic",
        }
        scheduler_name = scheduler_mapping.get(scheduler, scheduler) if scheduler else None

        # Build output lines
        prompt_line = prompt if prompt else ""
        if add_loras_to_prompt and loras_text:
            prompt_line = f"{prompt_line}\n{loras_text}" if prompt_line else loras_text
        lines = [prompt_line] if prompt_line else [""]
        if negative_prompt:
            lines.append(f"Negative prompt: {negative_prompt}")

        params: list[str] = []
        if steps is not None:
            params.append(f"Steps: {steps}")
        if sampler_name:
            if scheduler_name:
                params.append(f"Sampler: {sampler_name} {scheduler_name}")
            else:
                params.append(f"Sampler: {sampler_name}")
        if cfg is not None:
            params.append(f"CFG scale: {cfg}")
        if seed is not None:
            params.append(f"Seed: {seed}")
        if size:
            params.append(f"Size: {size}")
        if clip_skip is not None:
            try:
                params.append(f"Clip skip: {abs(int(clip_skip))}")
            except (ValueError, TypeError):
                pass
        additional_data = metadata_dict.get("additional_data", "")
        if additional_data:
            params.append(additional_data)
        if ckpt_hash:
            params.append(f"Model hash: {ckpt_hash[:10].upper()}")
        if ckpt_display_name:
            params.append(f"Model: {ckpt_display_name}")
        if hashes:
            params.append(f"Hashes: {json.dumps(hashes, separators=(',', ':'))}")
        params.append("Version: ComfyUI")
        if civitai_resources:
            params.append(
                f"Civitai resources: {json.dumps(civitai_resources, separators=(',', ':'))}"
            )

        lines.append(", ".join(params))
        return "\n".join(lines)

    # credit to nkchocoai
    # Add format_filename method to handle pattern substitution
    def format_filename(self, filename, metadata_dict):
        """Format filename with metadata values"""
        if not metadata_dict:
            return filename

        result = re.findall(self.pattern_format, filename)
        for segment in result:
            parts = segment.replace("%", "").split(":")
            key = parts[0]

            if key == "seed" and "seed" in metadata_dict:
                seed_value = metadata_dict.get("seed")
                if seed_value is not None:
                    filename = filename.replace(segment, str(seed_value))
                else:
                    # Fallback if seed was not captured by metadata collector
                    filename = filename.replace(segment, "0")
            elif key == "width" and "size" in metadata_dict:
                size = metadata_dict.get("size", "x")
                w = size.split("x")[0] if isinstance(size, str) else size[0]
                filename = filename.replace(segment, str(w))
            elif key == "height" and "size" in metadata_dict:
                size = metadata_dict.get("size", "x")
                h = size.split("x")[1] if isinstance(size, str) else size[1]
                filename = filename.replace(segment, str(h))
            elif key == "pprompt" and "prompt" in metadata_dict:
                prompt = metadata_dict.get("prompt", "").replace("\n", " ")
                prompt = sanitize_folder_name(prompt)
                if len(parts) >= 2:
                    length = int(parts[1])
                    prompt = prompt[:length]
                filename = filename.replace(segment, prompt.strip())
            elif key == "nprompt" and "negative_prompt" in metadata_dict:
                prompt = metadata_dict.get("negative_prompt", "").replace("\n", " ")
                prompt = sanitize_folder_name(prompt)
                if len(parts) >= 2:
                    length = int(parts[1])
                    prompt = prompt[:length]
                filename = filename.replace(segment, prompt.strip())
            elif key == "model":
                model_value = metadata_dict.get("checkpoint")
                if isinstance(model_value, (bytes, os.PathLike)):
                    model_value = str(model_value)

                if not isinstance(model_value, str) or not model_value:
                    model = "model_unavailable"
                else:
                    model = os.path.splitext(os.path.basename(model_value))[0]
                    model = sanitize_folder_name(model)
                if len(parts) >= 2:
                    length = int(parts[1])
                    model = model[:length]
                filename = filename.replace(segment, model)
            elif key == "date":
                from datetime import datetime

                now = datetime.now()
                date_table = {
                    "yyyy": f"{now.year:04d}",
                    "yy": f"{now.year % 100:02d}",
                    "MM": f"{now.month:02d}",
                    "dd": f"{now.day:02d}",
                    "hh": f"{now.hour:02d}",
                    "mm": f"{now.minute:02d}",
                    "ss": f"{now.second:02d}",
                }
                if len(parts) >= 2:
                    date_format = parts[1]
                    for k, v in date_table.items():
                        date_format = date_format.replace(k, v)
                    filename = filename.replace(segment, date_format)
                else:
                    date_format = "yyyyMMddhhmmss"
                    for k, v in date_table.items():
                        date_format = date_format.replace(k, v)
                    filename = filename.replace(segment, date_format)

        return filename

    @staticmethod
    def _get_cached_model_by_name(scanner, name):
        cache = getattr(scanner, "_cache", None)
        if cache is None or not name:
            return None

        candidates = [
            name,
            os.path.basename(name),
            os.path.splitext(os.path.basename(name))[0],
        ]
        for model in getattr(cache, "raw_data", []):
            file_name = model.get("file_name")
            if file_name in candidates:
                return model
        return None

    def _build_recipe_loras(self, recipe_scanner, lora_stack):
        lora_matches = re.findall(r"<lora:([^:]+):([^>]+)>", lora_stack or "")
        lora_scanner = getattr(recipe_scanner, "_lora_scanner", None)
        loras_data = []
        base_model_counts = {}

        for name, strength in lora_matches:
            lora_info = self._get_cached_model_by_name(lora_scanner, name)
            civitai = (lora_info or {}).get("civitai") or {}
            civitai_model = civitai.get("model") or {}
            try:
                parsed_strength = float(strength)
            except (TypeError, ValueError):
                parsed_strength = 1.0

            loras_data.append(
                {
                    "file_name": name,
                    "strength": parsed_strength,
                    "hash": ((lora_info or {}).get("sha256") or "").lower(),
                    "modelVersionId": civitai.get("id", 0),
                    "modelName": civitai_model.get("name", name) if lora_info else "",
                    "modelVersionName": civitai.get("name", "") if lora_info else "",
                    "isDeleted": False,
                    "exclude": False,
                }
            )

            base_model = (lora_info or {}).get("base_model")
            if base_model:
                base_model_counts[base_model] = base_model_counts.get(base_model, 0) + 1

        return lora_matches, loras_data, base_model_counts

    def _build_recipe_checkpoint(self, recipe_scanner, checkpoint_raw):
        if not isinstance(checkpoint_raw, str) or not checkpoint_raw.strip():
            return None

        checkpoint_name = checkpoint_raw.strip()
        file_name = os.path.splitext(os.path.basename(checkpoint_name))[0]
        checkpoint_scanner = getattr(recipe_scanner, "_checkpoint_scanner", None)
        checkpoint_info = self._get_cached_model_by_name(
            checkpoint_scanner, checkpoint_name
        )

        if not checkpoint_info:
            return {
                "type": "checkpoint",
                "name": checkpoint_name,
                "file_name": file_name,
                "hash": self.get_checkpoint_hash(checkpoint_name) or "",
            }

        civitai = checkpoint_info.get("civitai") or {}
        civitai_model = civitai.get("model") or {}
        file_path = checkpoint_info.get("file_path") or checkpoint_info.get("path") or ""
        cached_file_name = (
            checkpoint_info.get("file_name")
            or (os.path.splitext(os.path.basename(file_path))[0] if file_path else "")
            or file_name
        )

        return {
            "type": "checkpoint",
            "modelId": civitai_model.get("id", 0),
            "modelVersionId": civitai.get("id", 0),
            "name": civitai_model.get("name")
            or checkpoint_info.get("model_name")
            or checkpoint_name,
            "version": civitai.get("name", ""),
            "hash": (
                checkpoint_info.get("sha256") or checkpoint_info.get("hash") or ""
            ).lower(),
            "file_name": cached_file_name,
            "modelName": civitai_model.get("name", ""),
            "modelVersionName": civitai.get("name", ""),
            "baseModel": checkpoint_info.get("base_model")
            or civitai.get("baseModel", ""),
        }

    @staticmethod
    def _derive_recipe_name(lora_matches):
        recipe_name_parts = [
            f"{name.strip()}-{float(strength):.2f}" for name, strength in lora_matches[:3]
        ]
        return "_".join(recipe_name_parts) or "recipe"

    @staticmethod
    def _sync_recipe_cache(recipe_scanner, recipe_data, json_path):
        cache = getattr(recipe_scanner, "_cache", None)
        if cache is not None:
            cache.raw_data.append(recipe_data)
            cache.sorted_by_name = sorted(
                cache.raw_data, key=lambda item: item.get("title", "").lower()
            )
            cache.sorted_by_date = sorted(
                cache.raw_data,
                key=lambda item: (
                    item.get("modified", item.get("created_date", 0)),
                    item.get("file_path", ""),
                ),
                reverse=True,
            )
            recipe_scanner._update_folder_metadata(cache)
            recipe_scanner._update_fts_index_for_recipe(recipe_data, "add")

        recipe_id = str(recipe_data.get("id", ""))
        if recipe_id:
            recipe_scanner._json_path_map[recipe_id] = json_path
        persistent_cache = getattr(recipe_scanner, "_persistent_cache", None)
        if persistent_cache:
            persistent_cache.update_recipe(recipe_data, json_path)

    def _save_image_as_recipe(self, file_path, metadata_dict):
        if not metadata_dict:
            raise ValueError("No generation metadata found")

        recipe_scanner = ServiceRegistry.get_service_sync("recipe_scanner")
        if recipe_scanner is None:
            raise RuntimeError("Recipe scanner unavailable")

        recipes_dir = recipe_scanner.recipes_dir
        if not recipes_dir:
            raise RuntimeError("Recipes directory unavailable")
        os.makedirs(recipes_dir, exist_ok=True)

        recipe_id = str(uuid.uuid4())
        optimized_image, extension = ExifUtils.optimize_image(
            image_data=file_path,
            target_width=CARD_PREVIEW_WIDTH,
            format="webp",
            quality=85,
            preserve_metadata=True,
        )
        image_path = os.path.normpath(os.path.join(recipes_dir, f"{recipe_id}{extension}"))
        with open(image_path, "wb") as file_obj:
            file_obj.write(optimized_image)

        lora_stack = metadata_dict.get("loras", "")
        lora_matches, loras_data, base_model_counts = self._build_recipe_loras(
            recipe_scanner, lora_stack
        )
        checkpoint_entry = self._build_recipe_checkpoint(
            recipe_scanner, metadata_dict.get("checkpoint")
        )
        most_common_base_model = (
            max(base_model_counts.items(), key=lambda item: item[1])[0]
            if base_model_counts
            else ""
        )
        current_time = time.time()
        recipe_data = {
            "id": recipe_id,
            "file_path": image_path,
            "title": self._derive_recipe_name(lora_matches),
            "modified": current_time,
            "created_date": current_time,
            "base_model": most_common_base_model
            or (checkpoint_entry or {}).get("baseModel", ""),
            "loras": loras_data,
            "gen_params": {
                key: value
                for key, value in metadata_dict.items()
                if key not in ["checkpoint", "loras"]
            },
            "loras_stack": lora_stack,
            "fingerprint": calculate_recipe_fingerprint(loras_data),
        }
        if checkpoint_entry:
            recipe_data["checkpoint"] = checkpoint_entry

        # The recipe image is the WebP produced above from the output file;
        # reuse the same metadata extraction to record workflow presence.
        try:
            metadata = ExifUtils._load_structured_metadata(image_path)
            recipe_data["has_workflow"] = bool(metadata.get("workflow"))
        except Exception:
            recipe_data["has_workflow"] = False

        json_path = os.path.normpath(
            os.path.join(recipes_dir, f"{recipe_id}.recipe.json")
        )
        with open(json_path, "w", encoding="utf-8") as file_obj:
            json.dump(recipe_data, file_obj, indent=4, ensure_ascii=False)

        ExifUtils.append_recipe_metadata(image_path, recipe_data)
        self._sync_recipe_cache(recipe_scanner, recipe_data, json_path)

    def save_images(
        self,
        images,
        filename_prefix,
        file_format,
        id,
        prompt=None,
        extra_pnginfo=None,
        lossless_webp=True,
        quality=100,
        webp_method=6,
        jpeg_subsampling=0,
        embed_workflow=False,
        save_with_metadata=True,
        add_counter_to_filename=True,
        save_as_recipe=False,
        add_loras_to_prompt=False,
    ):
        """Save images with metadata"""
        results = []

        # Get metadata using the metadata collector
        raw_metadata = get_metadata()
        metadata_dict = MetadataProcessor.to_dict(raw_metadata, id)

        metadata = self.format_metadata(metadata_dict, add_loras_to_prompt)

        # Process filename_prefix with pattern substitution
        filename_prefix = self.format_filename(filename_prefix, metadata_dict)

        # Get initial save path info once for the batch
        full_output_folder, filename, counter, subfolder, processed_prefix = (
            folder_paths.get_save_image_path(
                filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
            )
        )

        # Create directory if it doesn't exist
        if not os.path.exists(full_output_folder):
            os.makedirs(full_output_folder, exist_ok=True)

        # Process each image with incrementing counter
        for i, image in enumerate(images):
            # Convert the tensor image to numpy array
            img = 255.0 * image.cpu().numpy()
            img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))

            # Generate filename with counter if needed
            base_filename = filename.replace("%batch_num%", str(i))
            if add_counter_to_filename:
                # Use counter + i to ensure unique filenames for all images in batch
                current_counter = counter + i
                base_filename += f"_{current_counter:05}_"

            # Set file extension and prepare saving parameters
            file: str
            save_kwargs: Dict[str, Any]
            pnginfo: Optional[PngImagePlugin.PngInfo] = None
            if file_format == "png":
                file = base_filename + ".png"
                file_extension = ".png"
                # Remove "optimize": True to match built-in node behavior
                save_kwargs = {"compress_level": self.compress_level}
                pnginfo = PngImagePlugin.PngInfo()
            elif file_format == "jpeg":
                file = base_filename + ".jpg"
                file_extension = ".jpg"
                save_kwargs = {"quality": quality, "optimize": True, "subsampling": jpeg_subsampling}
            elif file_format == "webp":
                file = base_filename + ".webp"
                file_extension = ".webp"
                save_kwargs = {
                    "quality": quality,
                    "lossless": lossless_webp,
                    "method": webp_method,
                }
            else:
                raise ValueError(f"Unsupported file format: {file_format}")

            # Full save path
            file_path = os.path.join(full_output_folder, file)

            # Save the image with metadata
            try:
                if file_format == "png":
                    assert pnginfo is not None
                    if save_with_metadata and metadata:
                        pnginfo.add_text("parameters", metadata)
                    if embed_workflow and extra_pnginfo is not None:
                        workflow_json = json.dumps(extra_pnginfo["workflow"])
                        pnginfo.add_text("workflow", workflow_json)
                    save_kwargs["pnginfo"] = pnginfo
                    img.save(file_path, format="PNG", **save_kwargs)
                elif file_format == "jpeg":
                    # For JPEG, use piexif
                    if save_with_metadata and metadata:
                        try:
                            exif_dict = {
                                "Exif": {
                                    piexif.ExifIFD.UserComment: b"UNICODE\0"
                                    + metadata.encode("utf-16be")
                                }
                            }
                            exif_bytes = piexif.dump(exif_dict)
                            save_kwargs["exif"] = exif_bytes
                        except Exception as e:
                            logger.error(f"Error adding EXIF data: {e}")
                    img.save(file_path, format="JPEG", **save_kwargs)
                elif file_format == "webp":
                    try:
                        # For WebP, use piexif for metadata
                        exif_dict = {}

                        if save_with_metadata and metadata:
                            exif_dict["Exif"] = {
                                piexif.ExifIFD.UserComment: b"UNICODE\0"
                                + metadata.encode("utf-16be")
                            }

                        # Add workflow if needed
                        if embed_workflow and extra_pnginfo is not None:
                            workflow_json = json.dumps(extra_pnginfo["workflow"])
                            exif_dict["0th"] = {
                                piexif.ImageIFD.ImageDescription: "Workflow:"
                                + workflow_json
                            }

                        exif_bytes = piexif.dump(exif_dict)
                        save_kwargs["exif"] = exif_bytes
                    except Exception as e:
                        logger.error(f"Error adding EXIF data: {e}")

                    img.save(file_path, format="WEBP", **save_kwargs)

                if save_as_recipe:
                    try:
                        self._save_image_as_recipe(file_path, metadata_dict)
                    except Exception as e:
                        logger.warning(
                            "Failed to save image as recipe: %s", e, exc_info=True
                        )

                results.append(
                    {"filename": file, "subfolder": subfolder, "type": self.type}
                )

            except Exception as e:
                logger.error(f"Error saving image: {e}")

        return results

    def process_image(
        self,
        images,
        id,
        filename_prefix="ComfyUI",
        file_format="png",
        prompt=None,
        extra_pnginfo=None,
        lossless_webp=True,
        quality=100,
        webp_method=6,
        jpeg_subsampling=0,
        embed_workflow=False,
        save_with_metadata=True,
        add_counter_to_filename=True,
        save_as_recipe=False,
        add_loras_to_prompt=False,
    ):
        """Process and save image with metadata"""
        # Make sure the output directory exists
        os.makedirs(self.output_dir, exist_ok=True)

        # If images is already a list or array of images, do nothing; otherwise, convert to list
        if isinstance(images, (list, np.ndarray)):
            pass
        else:
            # Ensure images is always a list of images
            if len(images.shape) == 3:  # Single image (height, width, channels)
                images = [images]
            else:  # Multiple images (batch, height, width, channels)
                images = [img for img in images]

        # Save all images
        results = self.save_images(
            images,
            filename_prefix,
            file_format,
            id,
            prompt,
            extra_pnginfo,
            lossless_webp,
            quality,
            webp_method,
            jpeg_subsampling,
            embed_workflow,
            save_with_metadata,
            add_counter_to_filename,
            save_as_recipe,
            add_loras_to_prompt,
        )

        return {
            "result": (images,),
            "ui": {"images": results},
        }
