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

Outputs a JSON mapping of filename → {"model": "...", "loras": ["...", ...]}.
"""

import json
import os
import sys
from pathlib import Path

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"}


def extract_ai_meta(filepath):
    """Return (model, loras) or (None, []) if no ComfyUI metadata found."""
    try:
        img = Image.open(filepath)
    except Exception:
        return None, []

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

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

    model = None
    loras = []

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

        if cls == "CheckpointLoaderSimple":
            ckpt = inputs.get("ckpt_name", "")
            if ckpt:
                model = ckpt.rsplit(".", 1)[0]

        if "lora" in cls.lower():
            lora = inputs.get("lora_name", "")
            if lora:
                loras.append(lora.rsplit(".", 1)[0])

    loras.sort()
    return model, loras


def main():
    target = sys.argv[1] if len(sys.argv) > 1 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 = {}

    for entry in sorted(os.listdir(target)):
        fpath = os.path.join(target, entry)
        if not os.path.isfile(fpath):
            continue
        ext = os.path.splitext(entry)[1].lower()
        if ext not in IMAGE_EXTS:
            continue

        model, loras = extract_ai_meta(fpath)
        if model or loras:
            meta[entry] = {"model": model, "loras": loras}

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


if __name__ == "__main__":
    main()