#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026 DiogenOS
#
# dio-img-prompts: a minimal viewer that lists the generation prompts
# embedded in PNG image metadata (ComfyUI-style "prompt" graphs and the
# "parameters" field written by other tools). Standard library only; no
# image-processing or machine-learning dependencies.

"""Graphical front end for recovering prompts from a directory of images."""

import os
import sys
import json
import zlib
import struct

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


def _png_text_fields(path):
    """Return a dict of textual key/value pairs stored in a PNG's chunks."""
    fields = {}
    with open(path, "rb") as handle:
        if handle.read(8) != PNG_SIGNATURE:
            return fields
        while True:
            header = handle.read(8)
            if len(header) < 8:
                break
            length, chunk_type = struct.unpack(">I4s", header)
            data = handle.read(length)
            handle.read(4)  # trailing CRC, not validated
            if chunk_type == b"tEXt":
                key, _, value = data.partition(b"\x00")
                fields[key.decode("latin1")] = value.decode("latin1")
            elif chunk_type == b"iTXt":
                key, _, rest = data.partition(b"\x00")
                if len(rest) < 2:
                    continue
                compressed = rest[0]
                rest = rest[2:]
                _, _, rest = rest.partition(b"\x00")  # language tag
                _, _, text = rest.partition(b"\x00")  # translated keyword
                try:
                    raw = zlib.decompress(text) if compressed == 1 else text
                    fields[key.decode("latin1")] = raw.decode("utf-8", "replace")
                except Exception:
                    fields[key.decode("latin1")] = text.decode("utf-8", "replace")
            elif chunk_type == b"IEND":
                break
    return fields


def _prompts_from_graph(graph):
    """Pull positive and negative prompt strings out of a ComfyUI graph."""
    def text_of(node_id):
        node = graph.get(str(node_id))
        if isinstance(node, dict):
            value = node.get("inputs", {}).get("text")
            if isinstance(value, str):
                return value.strip()
        return None

    positives, negatives = [], []
    sampler_seen = False
    for node in graph.values():
        if not isinstance(node, dict):
            continue
        inputs = node.get("inputs", {})
        if isinstance(inputs, dict) and "positive" in inputs and "negative" in inputs:
            sampler_seen = True
            for role, bucket in (("positive", positives), ("negative", negatives)):
                link = inputs.get(role)
                if isinstance(link, list) and link:
                    text = text_of(link[0])
                    if text:
                        bucket.append(text)

    if not sampler_seen:
        # No sampler resolved the roles; fall back to every text encoder node.
        for node in graph.values():
            if isinstance(node, dict) and "CLIPTextEncode" in str(node.get("class_type", "")):
                value = node.get("inputs", {}).get("text")
                if isinstance(value, str) and value.strip():
                    positives.append(value.strip())
    return positives, negatives


# Mapping of graph input keys to the label shown for them. The first scalar
# value found for each label wins, so loader/sampler nodes populate it once.
SETTING_KEYS = (
    ("ckpt_name", "model"),
    ("unet_name", "model"),
    ("model_name", "model"),
    ("vae_name", "vae"),
    ("seed", "seed"),
    ("noise_seed", "seed"),
    ("steps", "steps"),
    ("cfg", "cfg"),
    ("guidance", "guidance"),
    ("sampler_name", "sampler"),
    ("scheduler", "scheduler"),
    ("denoise", "denoise"),
    ("width", "width"),
    ("height", "height"),
)

SETTING_ORDER = ("model", "vae", "seed", "steps", "cfg", "guidance",
                 "sampler", "scheduler", "denoise")


def _settings_from_graph(graph):
    """Collect generation settings (seed, steps, model, ...) from a graph."""
    found = {}
    for node in graph.values():
        if not isinstance(node, dict):
            continue
        inputs = node.get("inputs", {})
        if not isinstance(inputs, dict):
            continue
        for key, label in SETTING_KEYS:
            if label in found or key not in inputs:
                continue
            value = inputs[key]
            if isinstance(value, (str, int, float, bool)):
                found[label] = value
    return found


def _settings_from_text(body):
    """Parse a plain-text "parameters" settings line (e.g. Steps: 20, ...)."""
    found = {}
    aliases = {
        "model": "model", "sampler": "sampler", "scheduler": "scheduler",
        "seed": "seed", "steps": "steps", "cfg scale": "cfg",
        "size": "size", "guidance": "guidance",
    }
    line = ""
    for candidate in body.splitlines():
        if "Steps:" in candidate or "Sampler:" in candidate:
            line = candidate
            break
    for segment in line.split(","):
        if ":" not in segment:
            continue
        key, _, value = segment.partition(":")
        label = aliases.get(key.strip().lower())
        if label:
            found[label] = value.strip()
    return found


def format_settings(found):
    """Render a settings dict as a single compact line."""
    if not found:
        return ""
    parts = []
    for label in SETTING_ORDER:
        if label in found:
            value = found[label]
            if label == "model":
                value = os.path.basename(str(value))
            parts.append("%s: %s" % (label, value))
    if "size" in found:
        parts.append("size: %s" % found["size"])
    elif "width" in found and "height" in found:
        parts.append("size: %sx%s" % (found["width"], found["height"]))
    return " | ".join(parts)


def extract_prompts(path):
    """Return (positives, negatives, settings) for a single image file."""
    fields = _png_text_fields(path)
    raw = fields.get("prompt") or fields.get("Comment")
    if raw:
        try:
            graph = json.loads(raw)
        except Exception:
            graph = None
        if isinstance(graph, dict):
            positives, negatives = _prompts_from_graph(graph)
            if positives or negatives:
                return positives, negatives, _settings_from_graph(graph)
    # Plain-text "parameters" field (no JSON graph): treat the whole value
    # as the prompt, splitting off a trailing negative section when present.
    if "parameters" in fields:
        body = fields["parameters"]
        settings = _settings_from_text(body)
        marker = "Negative prompt:"
        if marker in body:
            head, _, tail = body.partition(marker)
            return [head.strip()], [tail.strip().split("\n", 1)[0].strip()], settings
        head = body.split("Steps:", 1)[0]
        return [head.strip()], [], settings
    return [], [], {}


def scan_directory(root):
    """Walk a directory tree and collect prompts from every PNG found."""
    results = []
    scanned = 0
    for dirpath, _dirs, files in os.walk(root):
        for name in sorted(files):
            if not name.lower().endswith(".png"):
                continue
            full = os.path.join(dirpath, name)
            scanned += 1
            try:
                positives, negatives, settings = extract_prompts(full)
            except Exception:
                positives, negatives, settings = [], [], {}
            if positives or negatives:
                results.append((full, positives, negatives, settings))
    return scanned, results


def render(results, show_path):
    """Format results for display: per-image when show_path, else deduplicated."""
    lines = []
    if show_path:
        for path, positives, negatives, settings in results:
            lines.append(path)
            for text in positives:
                lines.append("  + " + text.replace("\n", "\n    "))
            for text in negatives:
                lines.append("  - " + text.replace("\n", "\n    "))
            line = format_settings(settings)
            if line:
                lines.append("  . " + line)
            lines.append("")
    else:
        seen = set()
        index = 0
        for _path, positives, _negatives, settings in results:
            for text in positives:
                if text in seen:
                    continue
                seen.add(text)
                index += 1
                lines.append("%d. %s" % (index, text))
                line = format_settings(settings)
                if line:
                    lines.append("   " + line)
                lines.append("")
    body = "\n".join(lines).strip()
    return body + "\n" if body else ""


class Application:
    def __init__(self, root, initial_dir=None):
        self.root = root
        root.title("Image Prompt Finder")
        root.geometry("840x620")
        root.minsize(560, 400)

        self.directory = tk.StringVar(value=initial_dir or os.path.expanduser("~"))
        self.show_path = tk.BooleanVar(value=False)
        self.status = tk.StringVar(value="Enter a directory and select Find prompts.")

        frame = ttk.Frame(root, padding=10)
        frame.pack(fill="both", expand=True)
        frame.columnconfigure(1, weight=1)

        ttk.Label(frame, text="Directory:").grid(row=0, column=0, sticky="w")
        entry = ttk.Entry(frame, textvariable=self.directory)
        entry.grid(row=0, column=1, sticky="ew", padx=6)
        entry.bind("<Return>", lambda _event: self.find())
        ttk.Button(frame, text="Browse", command=self.browse).grid(row=0, column=2)

        options = ttk.Frame(frame)
        options.grid(row=1, column=0, columnspan=3, sticky="w", pady=(8, 4))
        ttk.Checkbutton(
            options,
            text="Show image path with each prompt",
            variable=self.show_path,
        ).pack(side="left")

        buttons = ttk.Frame(frame)
        buttons.grid(row=2, column=0, columnspan=3, sticky="w", pady=(0, 6))
        ttk.Button(buttons, text="Find prompts", command=self.find).pack(side="left")
        ttk.Button(buttons, text="Copy all", command=self.copy_all).pack(side="left", padx=6)
        ttk.Label(buttons, textvariable=self.status).pack(side="left", padx=8)

        self.output = scrolledtext.ScrolledText(frame, wrap="word", undo=False)
        self.output.grid(row=3, column=0, columnspan=3, sticky="nsew")
        frame.rowconfigure(3, weight=1)
        self.output.bind("<Control-a>", self._select_all)
        self.output.bind("<Control-A>", self._select_all)

    def browse(self):
        chosen = filedialog.askdirectory(initialdir=self.directory.get() or os.path.expanduser("~"))
        if chosen:
            self.directory.set(chosen)

    def find(self):
        target = self.directory.get().strip()
        if not target:
            messagebox.showwarning("Image Prompt Finder", "Enter a directory to scan.")
            return
        target = os.path.expanduser(target)
        if not os.path.isdir(target):
            messagebox.showerror("Image Prompt Finder", "Not a directory:\n%s" % target)
            return
        self.status.set("Scanning ...")
        self.root.update_idletasks()
        scanned, results = scan_directory(target)
        text = render(results, self.show_path.get())
        self.output.delete("1.0", "end")
        self.output.insert("1.0", text)
        prompt_count = sum(len(p) for _f, p, _n, _s in results)
        self.status.set(
            "%d image(s) scanned, %d prompt-bearing, %d prompt(s) shown."
            % (scanned, len(results), prompt_count if self.show_path.get() else self._unique_count(results))
        )

    @staticmethod
    def _unique_count(results):
        seen = set()
        for _path, positives, _negatives, _settings in results:
            seen.update(positives)
        return len(seen)

    def copy_all(self):
        text = self.output.get("1.0", "end-1c")
        self.root.clipboard_clear()
        self.root.clipboard_append(text)
        self.status.set("Copied to clipboard.")

    def _select_all(self, _event):
        self.output.tag_add("sel", "1.0", "end-1c")
        return "break"


def main():
    # An optional directory argument prefills the field (and scans it at once),
    # so callers can open the finder pointed straight at an output folder.
    positional = [arg for arg in sys.argv[1:] if not arg.startswith("-")]
    initial = os.path.expanduser(positional[0]) if positional else None
    root = tk.Tk()
    app = Application(root, initial)
    if initial and os.path.isdir(initial):
        root.after(0, app.find)
    root.mainloop()


if __name__ == "__main__":
    main()
