#!/usr/bin/env python3
"""
rename-ai-snaps — Scan PNGs for AI prompt metadata and reorg into model/lora folders.

Usage:
  ./rename-ai-snaps [path] [options]

Scans for ComfyUI PNGs, reads embedded prompt metadata, extracts the
checkpoint model and LoRA names, and moves files into:
  {output}/{model}/{lora}/schmeeve-AI-{keywords}.png
"""

import json
import os
import re
import shutil
import sys
import time
from pathlib import Path

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

# ── stopwords and filter sets ──────────────────────────────────────────────

QUALITY_TAGS = {
    "score_6_up", "score_7_up", "score_8_up", "score_9",
    "score_6", "score_7", "score_8",
    "masterpiece", "best quality", "good quality", "normal quality",
    "high quality", "highly detailed", "very detailed", "extreme detail",
    "very_aesthetic", "absurdres", "8k", "4k",
    "photorealistic", "photograph",
    "depth of field", "solo focus", "cinematic",
    "newest", "amazing", "stunning",
}

QUALITY_WORDS = {
    "best", "good", "high", "top", "ultra", "super",
    "mega", "hyper", "extreme", "extra", "ultimate",
}

TECHNICAL_WORDS = {
    "detailed", "focus", "quality", "aesthetic", "realistic",
    "cinematic", "lighting", "rendering", "shading", "texture",
    "newest", "absurdres",
}

STOP_WORDS = {
    "the", "a", "an", "of", "in", "on", "at", "to", "for", "with",
    "and", "or", "is", "are", "was", "were", "be", "been", "being",
    "have", "has", "had", "do", "does", "did", "will", "would",
    "could", "should", "may", "might", "can", "shall", "this",
    "that", "these", "those", "it", "its", "by", "from", "as",
    "into", "through", "during", "before", "after", "above", "below",
    "between", "out", "off", "over", "under", "again", "further",
    "then", "once", "here", "there", "when", "where", "why", "how",
    "all", "each", "every", "both", "few", "more", "most", "other",
    "some", "such", "no", "nor", "not", "only", "own", "same", "so",
    "than", "too", "very", "just", "about", "up", "down",
    "make", "get", "set", "put", "take", "give", "show", "use",
    "like", "look", "see", "want", "need", "let", "close", "full",
    "add", "new", "one", "two", "five",
    "also", "well", "back", "still", "even", "much",
    "you", "your", "my", "me", "we", "our", "they", "them", "their",
}

GENERIC_WORDS = {
    "man", "men", "guy", "guys", "boy", "boys", "woman", "women",
    "girl", "girls", "people", "person", "human", "figure",
    "photo", "image", "picture", "shot", "view", "pose", "posing",
    "face", "head", "body", "skin", "hair", "eyes", "hand", "hands",
    "dark", "light", "bright", "color", "colour",
}

NEGATIVE_INDICATORS = {
    "deformed", "distorted", "disfigured", "poorly drawn", "bad anatomy",
    "extra digits", "missing digits", "extra limbs", "missing limbs",
    "ugly", "tiling", "low quality", "worst quality", "normal quality",
    "lowres", "monochrome", "grayscale", "text", "watermark",
    "branding", "border", "cropped", "signature", "username",
    "error", "mutation", "mutated", "out of frame", "duplicate", "cloned",
    "body out of frame", "bad hands", "bad face", "blurry",
}

# ── Helpers ────────────────────────────────────────────────────────────────

def spinner():
    chars = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
    i = 0
    while True:
        yield chars[i % len(chars)]
        i += 1


def is_negative_text(text):
    lower = text.lower()
    score = 0
    for ind in NEGATIVE_INDICATORS:
        if ind in lower:
            score += 1
    return score >= 2


def is_quality_only(text):
    lower = text.lower()
    words = re.findall(r"[a-z_]+", lower)
    if not words:
        return False
    meaningful = sum(1 for w in words if w not in QUALITY_TAGS and len(w) > 2)
    return meaningful == 0


# ── Metadata extraction from ComfyUI prompt JSON ───────────────────────────

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

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

MODEL_CLASS_TYPES = {
    "UNETLoader",
}


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)
    # Strip .safetensors extension
    if name.endswith(".safetensors"):
        name = name[:-len(".safetensors")]
    return name


def extract_model_name(data):
    """Return the checkpoint/UNET model name from the prompt JSON, or None."""
    model_name = None
    for node in data.values():
        inputs = node.get("inputs", {})
        class_type = node.get("class_type", "")
        if class_type in CHECKPOINT_CLASS_TYPES and "ckpt_name" in inputs:
            model_name = _strip_model_name(inputs["ckpt_name"])
            break
        if class_type in MODEL_CLASS_TYPES and "unet_name" in inputs:
            model_name = _strip_model_name(inputs["unet_name"])
            break
    return model_name if model_name else None


def extract_lora_names(data):
    """Return a sorted list of unique LoRA names from the prompt JSON."""
    names = []
    for node in data.values():
        inputs = node.get("inputs", {})
        class_type = node.get("class_type", "")
        if class_type in LORA_CLASS_TYPES and "lora_name" in inputs:
            names.append(_strip_model_name(inputs["lora_name"]))
        # Some custom nodes use 'lora' or similar — look for any field ending
        # in 'lora_name' or 'lora' that contains '.safetensors'
        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 and name not in names:
                    names.append(name)
    return sorted(set(names))


def extract_prompts(filepath):
    """Return (positive_prompt, negative_prompt) from a ComfyUI PNG."""
    try:
        img = Image.open(filepath)
    except Exception:
        return None, None

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

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

    # Gather all text fields from all nodes
    candidates = []
    for node in data.values():
        inputs = node.get("inputs", {})
        for field in ("text", "prompt", "positive", "negative", "value", "string"):
            val = inputs.get(field, "")
            if isinstance(val, str) and len(val.strip()) > 3:
                candidates.append(val.strip())

    # Second pass: resolve node references
    for node in data.values():
        inputs = node.get("inputs", {})
        for field in ("text", "prompt", "positive", "negative", "value", "string"):
            val = inputs.get(field)
            if isinstance(val, list) and len(val) == 2 and isinstance(val[0], str):
                ref_node = data.get(val[0], {})
                ref_inputs = ref_node.get("inputs", {})
                for rf in ("value", "text", "string"):
                    rv = ref_inputs.get(rf, "")
                    if isinstance(rv, str) and len(rv.strip()) > 3:
                        candidates.append(rv.strip())
                        break

    if not candidates:
        return None, None

    positives = [c for c in candidates if not is_negative_text(c)]
    negatives = [c for c in candidates if is_negative_text(c)]

    positives = [c for c in positives if not is_quality_only(c)]

    pos = max(positives, key=len) if positives else None
    neg = max(negatives, key=len) if negatives else None
    return pos, neg


def extract_all_candidates(text):
    """Extract all candidate keywords with position index for frequency scoring."""
    if not text:
        return []

    text_lower = text.lower()
    cleaned = re.sub(r"[\[\(\{][^\]\)\}]*[\]\)\}]", "", text_lower)
    segments = re.split(r"[,.;:!?]+", cleaned)

    seen = set()
    candidates = []
    position = 0

    for seg in segments:
        seg = seg.strip()
        if not seg or len(seg) < 4:
            continue

        words = re.findall(r"[a-zA-Z_]+", seg)
        good = []
        for w in words:
            wl = w.lower().strip("_")
            if len(wl) < 3:
                continue
            if wl in STOP_WORDS or wl in GENERIC_WORDS:
                continue
            if wl in QUALITY_TAGS or wl in QUALITY_WORDS:
                continue
            if wl in TECHNICAL_WORDS:
                continue
            if wl in NEGATIVE_INDICATORS:
                continue
            if wl.startswith("score") or wl.startswith("step"):
                continue
            if wl.isdigit():
                continue
            good.append(wl)

        if good:
            taken = 0
            for g in good:
                if g not in seen and taken < 2:
                    candidates.append((g, position))
                    seen.add(g)
                    taken += 1
                    position += 1

    return candidates


def select_best_keywords(candidates, freq_map, total_images, max_keywords=5):
    if not candidates:
        return []

    if total_images <= 1:
        return [kw for kw, _ in candidates[:max_keywords]]

    scored = []
    for kw, pos in candidates:
        freq = freq_map.get(kw, 1)
        rarity = 1.0 - ((freq - 1) / max(total_images - 1, 1))
        max_pos = max(len(candidates), 1)
        pos_bonus = (pos / max_pos) * 0.3
        score = rarity + pos_bonus
        scored.append((score, kw, pos))

    scored.sort(key=lambda x: (-x[0], x[2]))

    selected = []
    seen = set()
    for _, kw, _ in scored:
        if kw not in seen:
            selected.append(kw)
            seen.add(kw)
        if len(selected) >= max_keywords:
            break

    return selected


def keywords_to_name(keywords):
    if not keywords:
        return None

    desc = "-".join(keywords)
    desc = re.sub(r"[^a-z0-9-]", "-", desc)
    desc = re.sub(r"-+", "-", desc).strip("-")

    if len(desc) > 40:
        desc = desc[:40].rstrip("-")
        if "-" in desc:
            truncated = "-".join(desc.split("-")[:-1])
            if truncated and len(truncated) > 10:
                desc = truncated

    return f"schmeeve-AI-{desc}.png" if desc and len(desc) > 3 else None


def extract_metadata(filepath):
    """Return (model_name, lora_names, positive_prompt, negative_prompt)."""
    try:
        img = Image.open(filepath)
    except Exception:
        return None, None, None, None

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

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

    model = extract_model_name(data)
    loras = extract_lora_names(data)
    pos, neg = extract_prompts(filepath)
    return model, loras, pos, neg


def sanitize_dir_name(name):
    """Sanitize a string for use as a directory name."""
    name = name.strip().replace(" ", "_")
    name = re.sub(r"[^a-zA-Z0-9_.-]", "_", name)
    name = re.sub(r"_+", "_", name).strip("_")
    return name


# ── main ───────────────────────────────────────────────────────────────────

def nas_mount_base():
    """Return the NAS mount base path for this machine."""
    candidates = [
        "/mini.nas/miniShare1",
        os.path.expanduser("~/mnt/mini.nas/miniShare1"),
        os.path.expanduser("~/mnt/miniShare1"),
    ]
    for p in candidates:
        if os.path.isdir(p):
            return p
    return candidates[-1]


def main():
    import argparse

    DEFAULT_OUTPUT = os.path.join(nas_mount_base(), "Pictures", "Images", "AI")

    parser = argparse.ArgumentParser(
        description="Scan PNGs for AI prompt metadata and reorg into model/lora folders.",
    )
    parser.add_argument(
        "path", nargs="?", default=os.path.expanduser("~/Pictures"),
        help="Directory to scan for PNG files (default: ~/Pictures)",
    )
    parser.add_argument(
        "-o", "--output", default=DEFAULT_OUTPUT,
        help=f"Output root directory (default: {DEFAULT_OUTPUT})",
    )
    parser.add_argument(
        "-n", "--no-interactive", action="store_true",
        help="Auto-reorg without prompting",
    )
    parser.add_argument(
        "-r", "--recursive", action="store_true",
        help="Search directories recursively (default: off)",
    )
    parser.add_argument(
        "-p", "--preview", "--dry-run", action="store_true",
        dest="dry_run",
        help="Show proposed changes without making any (combine with -n for full listing)",
    )
    args = parser.parse_args()

    scan_dir = Path(args.path).expanduser().resolve()
    if not scan_dir.is_dir():
        print(f"Error: {scan_dir} is not a directory")
        sys.exit(1)

    output_root = Path(args.output).expanduser().resolve()

    # Gather PNGs
    glob_pattern = "**/*.png" if args.recursive else "*.png"
    pngs = sorted(scan_dir.glob(glob_pattern))

    if not pngs:
        print(f"  No PNG files found in {scan_dir}" +
              (" (recursive search)" if args.recursive else ""))
        return

    # ── Phase 1: Collect candidates and build frequency map ──
    sys.stdout.write("  Analyzing PNGs")
    sys.stdout.flush()
    spin = spinner()
    image_data = {}
    word_images = {}
    for p in pngs:
        sys.stdout.write(f"\r  {next(spin)} Analyzing PNGs")
        sys.stdout.flush()
        # Skip already-renamed files (schmeeve-AI-*)
        if p.name.startswith("schmeeve-AI-"):
            continue
        model, loras, pos, neg = extract_metadata(str(p))
        source = pos or neg
        if source:
            candidates = extract_all_candidates(source)
            if candidates:
                image_data[p] = {
                    "candidates": candidates,
                    "model": model,
                    "loras": loras,
                    "source": source,
                }
                for kw, _ in candidates:
                    if kw not in word_images:
                        word_images[kw] = set()
                    word_images[kw].add(p)

    freq_map = {kw: len(images) for kw, images in word_images.items()}
    total_images = len(image_data)

    # ── Phase 1b: Propose file paths ──
    raw_proposals = {}
    for p, info in image_data.items():
        candidates = info["candidates"]
        model = info["model"]
        loras = info["loras"]

        if total_images > 1:
            keywords = select_best_keywords(candidates, freq_map, total_images)
        else:
            keywords = [kw for kw, _ in candidates[:5]]
        name = keywords_to_name(keywords)
        if not name:
            continue

        # Determine folder structure
        model_dir = sanitize_dir_name(model) if model else "unknown"
        if loras:
            lora_dir = sanitize_dir_name("+".join(loras))
        else:
            lora_dir = "no-lora"

        rel_path = Path(model_dir) / lora_dir / name
        raw_proposals[p] = rel_path

    # Clear the spinner line
    sys.stdout.write("\r" + " " * 60 + "\r")
    sys.stdout.flush()

    # Deduplicate file names within their target directories
    proposals = {}
    used_paths = {}
    for p, rel_path in raw_proposals.items():
        new_path = rel_path
        counter = 0
        key = str(new_path)
        while key in used_paths:
            counter += 1
            ts = time.strftime("%m%d-%H%M%S", time.localtime(p.stat().st_mtime))
            stem = rel_path.stem
            ext = rel_path.suffix
            new_path = rel_path.with_name(f"{stem}-{ts}{'_' + str(counter) if counter > 1 else ''}{ext}")
            key = str(new_path)
        used_paths[key] = True
        proposals[p] = new_path

    # ── Phase 2: Display proposed changes ──
    if not proposals:
        print("  No renamable PNGs found.")
        return

    print(f"\n  {'DRY RUN' if args.dry_run else 'PROPOSED'} — "
          f"{len(proposals)} file(s) to organize into {output_root}/\n")

    for i, (old_path, rel_path) in enumerate(proposals.items(), 1):
        try:
            src_display = str(old_path.relative_to(scan_dir))
        except ValueError:
            src_display = str(old_path)
        if len(src_display) > 55:
            src_display = src_display[:25] + "…" + src_display[-27:]
        print(f"  {i:>3}. {src_display}")
        print(f"       → {output_root / rel_path}")

    print()

    if args.dry_run:
        print(f"  Dry-run complete. No files were changed.\n")
        return

    # ── Phase 3: Rename/move (interactive or auto) ──
    if args.no_interactive:
        moved = 0
        for old_path, rel_path in proposals.items():
            new_path = output_root / rel_path
            new_path.parent.mkdir(parents=True, exist_ok=True)
            if new_path.exists():
                stem = new_path.stem
                counter = 1
                while new_path.exists():
                    new_path = new_path.with_name(f"{stem}_{counter}{new_path.suffix}")
                    counter += 1
            try:
                shutil.move(old_path, new_path)
                moved += 1
            except (FileNotFoundError, OSError) as e:
                print(f"  SKIPPED ({e}): {old_path.name}")
        print(f"  Moved {moved} file(s)." if not args.dry_run else "")
    else:
        moved = 0
        skipped = 0
        items = list(proposals.items())
        i = 0
        while i < len(items):
            old_path, rel_path = items[i]
            new_path = output_root / rel_path
            new_path_display = str(rel_path)

            print(f"\n  [{i+1}/{len(items)}]")
            print(f"  From: {old_path.name}")
            print(f"  To:   {output_root / rel_path}")
            remaining = len(items) - i - 1
            rlabel = f"move all {remaining}" if remaining else ""
            sys.stdout.write(f"  [Enter]=move  [e]=edit  [s]=skip{'  [a]=' + rlabel if rlabel else ''}  [q]=quit: ")
            sys.stdout.flush()
            choice = sys.stdin.readline().strip().lower()

            if choice == "q":
                remaining = len(items) - i - 1
                if remaining:
                    print(f"  Skipping remaining {remaining} file(s).")
                break
            elif choice == "s":
                skipped += 1
                i += 1
                continue
            elif choice == "e":
                sys.stdout.write(f"  Edit name (will be 'schmeeve-AI-{rel_path.parent}/...'): ")
                sys.stdout.flush()
                custom = sys.stdin.readline().strip()
                if custom:
                    custom_desc = re.sub(r"[^a-z0-9-]", "-", custom.lower())
                    custom_desc = re.sub(r"-+", "-", custom_desc).strip("-")
                    if custom_desc:
                        new_name = f"schmeeve-AI-{custom_desc}.png"
                        rel_path = rel_path.with_name(new_name)
                    else:
                        print("  Invalid name, skipping.")
                        skipped += 1
                        i += 1
                        continue
                else:
                    skipped += 1
                    i += 1
                    continue
            elif choice == "":
                pass
            elif choice == "a":
                for j in range(i, len(items)):
                    p, rp = items[j]
                    np = output_root / rp
                    np.parent.mkdir(parents=True, exist_ok=True)
                    if np.exists():
                        stem = np.stem
                        counter = 1
                        while np.exists():
                            np = np.with_name(f"{stem}_{counter}{np.suffix}")
                            counter += 1
                    try:
                        shutil.move(p, np)
                    except (FileNotFoundError, OSError) as e:
                        print(f"  SKIPPED ({e}): {p.name}")
                moved += len(items) - i
                break
            else:
                print(f"  Unknown option '{choice}', skipping.")
                skipped += 1
                i += 1
                continue

            # Perform move
            new_path = output_root / rel_path
            new_path.parent.mkdir(parents=True, exist_ok=True)
            if new_path.exists():
                stem = new_path.stem
                counter = 1
                while new_path.exists():
                    new_path = new_path.with_name(f"{stem}_{counter}{new_path.suffix}")
                    counter += 1
                print(f"  (file existed, saved as {new_path.name})")
            try:
                shutil.move(old_path, new_path)
                moved += 1
            except (FileNotFoundError, OSError) as e:
                print(f"  SKIPPED ({e}): {old_path.name}")
            i += 1

        print(f"\n  Moved: {moved}  Skipped: {skipped}")

    # Check for JPEGs with AI metadata
    jpg_pattern = "**/*.jpg" if args.recursive else "*.jpg"
    remaining_ai = 0
    for f in scan_dir.glob(jpg_pattern):
        try:
            img = Image.open(f)
            if "prompt" in img.info:
                remaining_ai += 1
        except Exception:
            pass
    if remaining_ai:
        print(f"  Note: {remaining_ai} JPEG(s) with AI metadata found (not yet supported).")


if __name__ == "__main__":
    main()