#!/usr/bin/env python3
"""
rename-ai-snaps — Scan PNGs for AI prompt metadata and rename descriptively.

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

Scans for ComfyUI PNGs, reads embedded prompt metadata, extracts keywords,
and renames files IN PLACE to:
  schmeeve-AI-{keywords}.png

Renaming only — use organize-images to sort files into model/lora folders.
"""

import json
import os
import re
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",
}

# Generic ComfyUI output filenames (ComfyUI_00042_.png, ComfyUI_temp_….png)
GENERIC_RE = re.compile(r"^ComfyUI", re.IGNORECASE)

# ── 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 ───────────────────────────

def load_prompt_data(filepath):
    """Return parsed ComfyUI prompt JSON from a PNG, or None."""
    try:
        img = Image.open(filepath)
    except Exception:
        return None

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

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


def extract_prompts(data):
    """Return (positive_prompt, negative_prompt) from parsed prompt JSON."""
    # 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


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

def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Scan PNGs for AI prompt metadata and rename in place. "
                    "Does not move files — use organize-images for sorting.",
    )
    parser.add_argument(
        "path", nargs="?", default=os.path.expanduser("~/Pictures"),
        help="Directory to scan for PNG files (default: ~/Pictures)",
    )
    parser.add_argument(
        "-n", "--no-interactive", action="store_true",
        help="Auto-rename 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",
    )
    parser.add_argument(
        "--include-renamed", action="store_true",
        help="Also process files already named schmeeve-AI-* (default: skip them)",
    )
    parser.add_argument(
        "--only-generic", action="store_true",
        help="Only rename generic ComfyUI_* filenames",
    )
    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)

    # 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-*) unless requested
        if not args.include_renamed and p.name.startswith("schmeeve-AI-"):
            continue
        if args.only_generic and not GENERIC_RE.match(p.name):
            continue
        data = load_prompt_data(str(p))
        if not data:
            continue
        pos, neg = extract_prompts(data)
        source = pos or neg
        if source:
            candidates = extract_all_candidates(source)
            if candidates:
                image_data[p] = candidates
                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 new names (in place — same directory) ──
    raw_proposals = {}
    for p, candidates in image_data.items():
        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 or name == p.name:
            continue
        raw_proposals[p] = name

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

    # Deduplicate file names within their directories
    proposals = {}
    used_paths = {}
    for p, name in raw_proposals.items():
        new_path = p.with_name(name)
        key = str(new_path)
        counter = 0
        while key in used_paths:
            counter += 1
            ts = time.strftime("%m%d-%H%M%S", time.localtime(p.stat().st_mtime))
            stem = new_path.stem
            ext = new_path.suffix
            new_path = new_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 rename in place\n")

    for i, (old_path, new_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"       → {new_path.name}")

    print()

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

    def do_rename(old_path, new_path):
        """Rename with collision suffix. Returns final path or None on failure."""
        if old_path == new_path:
            return new_path
        if new_path.exists():
            stem = new_path.stem
            ext = new_path.suffix
            counter = 1
            while new_path.exists():
                new_path = new_path.with_name(f"{stem}_{counter}{ext}")
                counter += 1
            print(f"  (name existed, using {new_path.name})")
        try:
            old_path.rename(new_path)
            return new_path
        except (FileNotFoundError, OSError) as e:
            print(f"  SKIPPED ({e}): {old_path.name}")
            return None

    # ── Phase 3: Rename (interactive or auto) ──
    if args.no_interactive:
        renamed = 0
        for old_path, new_path in proposals.items():
            if do_rename(old_path, new_path):
                renamed += 1
        print(f"  Renamed {renamed} file(s).")
    else:
        renamed = 0
        skipped = 0
        items = list(proposals.items())
        i = 0
        while i < len(items):
            old_path, new_path = items[i]

            print(f"\n  [{i+1}/{len(items)}]")
            print(f"  From: {old_path.name}")
            print(f"  To:   {new_path.name}")
            remaining = len(items) - i - 1
            rlabel = f"rename all {remaining}" if remaining else ""
            sys.stdout.write(f"  [Enter]=rename  [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 keywords (name becomes 'schmeeve-AI-...'): ")
                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_path = old_path.with_name(f"schmeeve-AI-{custom_desc}.png")
                    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, np = items[j]
                    if do_rename(p, np):
                        renamed += 1
                break
            else:
                print(f"  Unknown option '{choice}', skipping.")
                skipped += 1
                i += 1
                continue

            if do_rename(old_path, new_path):
                renamed += 1
            i += 1

        print(f"\n  Renamed: {renamed}  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()