#!/usr/bin/env python3
"""Extract model and LoRA metadata from ComfyUI-generated images.

Outputs a JSON mapping of path (relative to target) →
  {"model": "...", "loras": ["...", ...], "dest": "{model_dir}/{lora_dir}"}

Only files with parseable ComfyUI 'prompt' metadata are included, so
presence in the output doubles as "is an AI-generated image" for
organize-images. 'dest' is the canonical folder path (sans leading "AI/"):
model/LoRA names are path-stripped and sanitized (spaces/specials → _),
LoRAs joined with '+', missing model → "unknown", no LoRAs → "no-lora".

Usage: extract-ai-meta [path] [-r|--recursive]
"""

import json
import os
import re
import sys

try:
    from PIL import Image
except ImportError:
    print("Error: Pillow (PIL) required. Install with: pip install Pillow", file=sys.stderr)
    sys.exit(1)

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}

CHECKPOINT_CLASS_TYPES = {
    "CheckpointLoaderSimple",
    "Checkpoint Loader with Name (Image Saver)",
    "CheckpointLoader",
}

MODEL_CLASS_TYPES = {
    "UNETLoader",
}

LORA_CLASS_TYPES = {
    "LoraLoader",
    "Power Lora Loader (rgthree)",
}

# Filesystem-safe cap for the joined LoRA combo directory name
MAX_LORA_DIR_LEN = 200


def _strip_model_name(raw):
    """Strip path prefix and .safetensors extension from a model/LoRA name."""
    name = raw.strip()
    # Handle backslash/forward-slash paths: "Pony\model.safetensors" or "Pony/model.safetensors"
    name = name.replace("\\", "/")
    name = os.path.basename(name)
    if name.endswith(".safetensors"):
        name = name[:-len(".safetensors")]
    return name


def sanitize_dir_name(name):
    """Canonical directory-name sanitizer (spaces/specials → _)."""
    name = name.strip().replace(" ", "_")
    name = re.sub(r"[^a-zA-Z0-9_.-]", "_", name)
    name = re.sub(r"_+", "_", name).strip("_")
    return name


def is_lora_connected(node_id, data):
    """Check if a LoRA node's output is consumed by any downstream node."""
    for nid, node in data.items():
        if nid == node_id:
            continue
        for val in node.get("inputs", {}).values():
            if isinstance(val, list) and len(val) == 2 and str(val[0]) == str(node_id):
                return True
    return False


def extract_ai_meta(filepath):
    """Return (model, loras, has_prompt).

    has_prompt is True whenever the file carries parseable ComfyUI 'prompt'
    metadata, even if no model/LoRA nodes were found in it.
    """
    try:
        img = Image.open(filepath)
    except Exception:
        return None, [], False

    if "prompt" not in img.info:
        return None, [], False

    try:
        data = json.loads(img.info["prompt"])
    except (json.JSONDecodeError, TypeError):
        return None, [], False

    model = None
    loras = []

    for nid, node in data.items():
        cls = node.get("class_type", "")
        inputs = node.get("inputs", {})

        if model is None:
            if cls in CHECKPOINT_CLASS_TYPES and inputs.get("ckpt_name"):
                model = _strip_model_name(inputs["ckpt_name"])
            elif cls in MODEL_CLASS_TYPES and inputs.get("unet_name"):
                model = _strip_model_name(inputs["unet_name"])

        if cls in LORA_CLASS_TYPES and inputs.get("lora_name"):
            # Only include LoRAs whose output is actually consumed
            if is_lora_connected(nid, data):
                loras.append(_strip_model_name(inputs["lora_name"]))
        else:
            # Some custom nodes use fields ending in 'lora'/'lora_name'
            for field, val in inputs.items():
                if isinstance(val, str) and val.endswith(".safetensors") and "lora" in field.lower():
                    name = _strip_model_name(val)
                    if name:
                        loras.append(name)

    loras = sorted(set(loras))
    return model, loras, True


def dest_for(model, loras):
    """Canonical relative destination dir (sans 'AI/' prefix)."""
    model_dir = sanitize_dir_name(model) if model else "unknown"
    if loras:
        lora_dir = "+".join(sanitize_dir_name(l) for l in loras if l)
        if len(lora_dir) > MAX_LORA_DIR_LEN:
            lora_dir = lora_dir[:MAX_LORA_DIR_LEN].rstrip("+_")
        if not lora_dir:
            lora_dir = "no-lora"
    else:
        lora_dir = "no-lora"
    return f"{model_dir}/{lora_dir}"


def iter_images(target, recursive):
    if recursive:
        for dirpath, _dirnames, filenames in os.walk(target):
            for fn in sorted(filenames):
                if os.path.splitext(fn)[1].lower() in IMAGE_EXTS:
                    yield os.path.join(dirpath, fn)
    else:
        for entry in sorted(os.listdir(target)):
            fpath = os.path.join(target, entry)
            if os.path.isfile(fpath) and os.path.splitext(entry)[1].lower() in IMAGE_EXTS:
                yield fpath


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("-")]
    recursive = any(a in ("-r", "--recursive") for a in sys.argv[1:])
    target = args[0] if args else os.path.expanduser("~/Pictures")
    target = os.path.abspath(target)

    if not os.path.isdir(target):
        print(f"Error: {target} is not a directory", file=sys.stderr)
        sys.exit(1)

    meta = {}
    scanned = 0
    for fpath in iter_images(target, recursive):
        scanned += 1
        if scanned % 200 == 0:
            print(f"\r  …scanned {scanned} images ({len(meta)} with AI metadata)",
                  end="", file=sys.stderr, flush=True)
        model, loras, has_prompt = extract_ai_meta(fpath)
        if not has_prompt:
            continue
        rel = os.path.relpath(fpath, target)
        meta[rel] = {
            "model": model,
            "loras": loras,
            "dest": dest_for(model, loras),
        }

    if scanned:
        print(f"\r  Scanned {scanned} images, {len(meta)} with AI metadata.      ",
              file=sys.stderr, flush=True)

    print(json.dumps(meta, indent=2))


if __name__ == "__main__":
    main()