"""
comfyui_meta.py
================
Extraction and readable structuring of the ComfyUI metadata embedded in
a PNG, MP3, FLAC or MP4/MOV file (workflow + prompt).

This module has no Gradio dependency: it can be tested standalone from
the command line:

    python3 comfyui_meta.py my_file.png

General logic
-------------
ComfyUI embeds two JSON blobs in the files it produces:

- "prompt"   : the "API" graph actually sent to the server for
               generation. Values here are already resolved (e.g. a
               dynamic prompt "{a|b}" already shows up picked as "a"
               or "b").
- "workflow" : the "UI" graph as displayed in ComfyUI (positions,
               titles, and most importantly "widgets_values" = what
               the user typed/set on each node, BEFORE resolution ->
               this is where the "{a|b|c}" template of a dynamic
               prompt node is found).

We cross-reference the two to obtain, for each prompt actually used:
  - the final (resolved) value that was really sent to the engine
  - the original dynamic template, if a wildcard/dynamic prompt node
    was detected upstream.
"""

from __future__ import annotations

import json
import re
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


# ---------------------------------------------------------------------------
# 1. Raw extraction of the JSON blobs, depending on file type
# ---------------------------------------------------------------------------

class MetadataNotFound(Exception):
    pass


def extract_raw_metadata(path: str | Path) -> dict[str, Any]:
    """Returns {'prompt': <dict|None>, 'workflow': <dict|None>} for a
    file. Raises MetadataNotFound if nothing could be extracted."""
    path = Path(path)
    suffix = path.suffix.lower()

    if suffix in {".png", ".webp"}:
        data = _extract_from_image(path)
    elif suffix in {".mp3", ".flac", ".m4a", ".ogg"}:
        data = _extract_from_audio(path)
    elif suffix in {".mp4", ".mov", ".webm", ".mkv"}:
        data = _extract_from_video(path)
    else:
        # last resort: try image, then audio
        data = _extract_from_image(path) or _extract_from_audio(path)

    if not data or not (data.get("prompt") or data.get("workflow")):
        raise MetadataNotFound(
            f"No ComfyUI metadata (prompt/workflow) found in {path.name}."
        )
    return data


def _try_json(text: Any) -> Any:
    if not isinstance(text, str):
        return None
    text = text.strip()
    if not text:
        return None
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return None


def _extract_from_image(path: Path) -> dict[str, Any] | None:
    try:
        from PIL import Image
    except ImportError:
        return None
    try:
        img = Image.open(path)
    except Exception:
        return None
    info = getattr(img, "info", {}) or {}
    result = {"prompt": _try_json(info.get("prompt")), "workflow": _try_json(info.get("workflow"))}

    if not result["prompt"] and not result["workflow"]:
        # PNG stores "prompt"/"workflow" as plain text chunks, which PIL
        # exposes directly through img.info (handled above). WEBP has no
        # such text-chunk mechanism, so ComfyUI instead embeds the same
        # JSON blobs inside EXIF tags, each value prefixed as
        # "<label>:<json>" (e.g. "prompt:{...}"). Look there next.
        result = _extract_from_exif(img, result)

    return result


def _extract_from_exif(img, result: dict[str, Any]) -> dict[str, Any]:
    exif = None
    try:
        exif = img.getexif()
    except Exception:
        exif = None

    if not exif:
        # Some Pillow/format combinations don't populate getexif()
        # automatically even though raw EXIF bytes are present — parse
        # those bytes directly as a fallback.
        raw = (getattr(img, "info", {}) or {}).get("exif")
        if raw:
            try:
                from PIL import Image as _Image
                exif = _Image.Exif()
                exif.load(raw)
            except Exception:
                exif = None

    if not exif:
        return result

    for tag_id in exif:
        try:
            value = exif[tag_id]
        except Exception:
            continue
        if isinstance(value, bytes):
            try:
                value = value.decode("utf-8", errors="ignore")
            except Exception:
                continue
        if not isinstance(value, str):
            continue

        label, sep, json_part = value.partition(":")
        candidate = json_part if sep else value
        parsed = _try_json(candidate)
        if parsed is None:
            continue

        label = label.strip().lower()
        if "prompt" in label and not result.get("prompt"):
            result["prompt"] = parsed
        elif "workflow" in label and not result.get("workflow"):
            result["workflow"] = parsed
        elif not result.get("prompt") and not result.get("workflow") and isinstance(parsed, dict):
            # Unlabeled (or differently labeled) but still a JSON object —
            # guess which one it is from its shape: a "workflow" graph has
            # top-level "nodes"/"links", the "prompt" API graph doesn't.
            if "nodes" in parsed and "links" in parsed:
                result["workflow"] = parsed
            else:
                result["prompt"] = parsed

    return result


def _extract_from_audio(path: Path) -> dict[str, Any] | None:
    try:
        import mutagen
    except ImportError:
        return None

    result: dict[str, Any] = {"prompt": None, "workflow": None}
    try:
        f = mutagen.File(path)
    except Exception:
        f = None

    if f is None:
        return result

    tags = getattr(f, "tags", None)
    if tags is None:
        return result

    # MP3 / ID3 (TXXX:prompt, TXXX:workflow)
    try:
        from mutagen.id3 import ID3
        if isinstance(tags, ID3) or path.suffix.lower() == ".mp3":
            for frame_key in list(tags.keys()):
                if frame_key.upper().startswith("TXXX:"):
                    label = frame_key.split(":", 1)[1].strip().lower()
                    frame = tags[frame_key]
                    value = frame.text[0] if getattr(frame, "text", None) else None
                    parsed = _try_json(value)
                    if parsed is None:
                        continue
                    if "prompt" in label:
                        result["prompt"] = parsed
                    elif "workflow" in label:
                        result["workflow"] = parsed
            if result["prompt"] or result["workflow"]:
                return result
    except ImportError:
        pass

    # Vorbis comments (FLAC/OGG) or generic "dict-like" tags (M4A, etc.)
    try:
        items = tags.items() if hasattr(tags, "items") else []
    except Exception:
        items = []
    for key, value in items:
        key_l = str(key).lower()
        raw = value[0] if isinstance(value, (list, tuple)) and value else value
        parsed = _try_json(raw)
        if parsed is None:
            continue
        if "prompt" in key_l and result["prompt"] is None:
            result["prompt"] = parsed
        elif "workflow" in key_l and result["workflow"] is None:
            result["workflow"] = parsed

    return result


def _extract_from_video(path: Path) -> dict[str, Any] | None:
    """Best effort: clips exported by VideoHelperSuite generally embed
    the workflow/prompt in container tags (comment, description...).
    We query ffprobe if available."""
    result: dict[str, Any] = {"prompt": None, "workflow": None}
    try:
        proc = subprocess.run(
            [
                "ffprobe", "-v", "quiet", "-print_format", "json",
                "-show_entries", "format_tags", str(path),
            ],
            capture_output=True, text=True, timeout=20,
        )
        info = json.loads(proc.stdout or "{}")
        tags = info.get("format", {}).get("tags", {})
    except Exception:
        tags = {}

    for key, value in tags.items():
        key_l = key.lower()
        parsed = _try_json(value)
        if parsed is None:
            continue
        if "prompt" in key_l and result["prompt"] is None:
            result["prompt"] = parsed
        elif "workflow" in key_l and result["workflow"] is None:
            result["workflow"] = parsed

    # If nothing found, some exports put everything (prompt+workflow) in
    # a single "comment" field as one object {"prompt":..,"workflow":..}
    if not result["prompt"] and not result["workflow"]:
        for key in ("comment", "description"):
            parsed = _try_json(tags.get(key))
            if isinstance(parsed, dict) and ("prompt" in parsed or "workflow" in parsed):
                result["prompt"] = parsed.get("prompt")
                result["workflow"] = parsed.get("workflow")
                break

    return result


# ---------------------------------------------------------------------------
# 2. Data model for the structured view
# ---------------------------------------------------------------------------

@dataclass
class ModelEntry:
    node: str
    kind: str          # "Checkpoint", "UNET", "VAE", "CLIP", "LoRA"
    name: str
    detail: str = ""


@dataclass
class ParamEntry:
    node: str
    param: str
    value: str


@dataclass
class PromptEntry:
    node: str
    role: str           # e.g. "Positive", "Tags", "Lyrics"
    resolved: str
    template: str | None = None   # raw { | } template, if found upstream
    source_node: str | None = None


@dataclass
class StructuredView:
    models: list[ModelEntry] = field(default_factory=list)
    params: list[ParamEntry] = field(default_factory=list)
    prompts: list[PromptEntry] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# 3. Categorization
# ---------------------------------------------------------------------------

_MODEL_CLASS_RE = re.compile(
    r"checkpoint|unetloader|vaeloader|cliploader|controlnet|ipadapter", re.I
)
_LORA_RE = re.compile(r"lora", re.I)
_PARAM_CLASS_RE = re.compile(
    r"ksampler|samplercustom|emptylatent|emptyacestep|modelsampling|latentaudio",
    re.I,
)
_ENCODE_CLASS_RE = re.compile(r"textencode|cliptextencode", re.I)
_DYNAMIC_CLASS_RE = re.compile(r"dynamic.?prompt|wildcard", re.I)
_BRACE_PATTERN_RE = re.compile(r"\{[^{}]*\|[^{}]*\}")

_TEXT_KEY_PRIORITY = ["Text", "text", "string", "STRING", "value",
                       "wildcard_text", "populated_text"]

_ROLE_LABELS = {
    "text": "Prompt",
    "positive": "Positive",
    "negative": "Negative",
    "tags": "Tags (style)",
    "lyrics": "Lyrics",
    "text_g": "Prompt (global)",
    "text_l": "Prompt (local)",
}


def _node_label(node_id: str, workflow_by_id: dict[str, dict], class_type: str) -> str:
    wf_node = workflow_by_id.get(node_id)
    title = wf_node.get("title") if wf_node else None
    if title:
        return f"{title} (#{node_id})"
    return f"{class_type} (#{node_id})"


def _flatten_strings(value: Any) -> list[str]:
    """Some nodes (e.g. ShowText|pysssss) store their widget as nested
    lists (value history). We flatten to recover all non-empty strings,
    in order."""
    out: list[str] = []
    if isinstance(value, str):
        if value.strip():
            out.append(value)
    elif isinstance(value, (list, tuple)):
        for item in value:
            out.extend(_flatten_strings(item))
    return out


def _first_string_widget(wf_node: dict | None) -> str | None:
    if not wf_node:
        return None
    strings = _flatten_strings(wf_node.get("widgets_values"))
    if not strings:
        return None
    # prefer the longest string (usually the real prompt, rather than a
    # short flag/mode string like "enable" or "randomize")
    return max(strings, key=len)


def _resolve_link_text(prompt_graph: dict, workflow_by_id: dict, link: list) -> tuple[str | None, str | None]:
    """Resolves the final value of a [node_id, output_index] link
    pointing at a text "source" node. Returns (value, source_node_id)."""
    src_id = str(link[0])
    src_node = prompt_graph.get(src_id, {})
    inputs = src_node.get("inputs", {})
    for key in _TEXT_KEY_PRIORITY:
        if key in inputs and isinstance(inputs[key], str) and inputs[key].strip():
            return inputs[key], src_id
    # fallback: last displayed/typed value on the source node (widget)
    fallback = _first_string_widget(workflow_by_id.get(src_id))
    if fallback:
        return fallback, src_id
    return None, src_id


def build_structured_view(raw: dict[str, Any]) -> StructuredView:
    view = StructuredView()
    prompt_graph: dict = raw.get("prompt") or {}
    workflow: dict = raw.get("workflow") or {}
    workflow_nodes = workflow.get("nodes") or []
    workflow_by_id = {str(n.get("id")): n for n in workflow_nodes}

    if not prompt_graph:
        view.warnings.append(
            "No 'prompt' (API) graph found — only 'workflow' info "
            "(if present) is shown."
        )

    # --- Model & LoRAs ---------------------------------------------
    for node_id, node in prompt_graph.items():
        class_type = node.get("class_type", "")
        inputs = node.get("inputs", {})
        label = _node_label(node_id, workflow_by_id, class_type)

        if _MODEL_CLASS_RE.search(class_type):
            for key in ("ckpt_name", "unet_name", "vae_name",
                        "clip_name", "clip_name1", "clip_name2", "clip_name3",
                        "control_net_name"):
                if key in inputs and isinstance(inputs[key], str):
                    kind = {
                        "ckpt_name": "Checkpoint",
                        "unet_name": "UNET",
                        "vae_name": "VAE",
                        "control_net_name": "ControlNet",
                    }.get(key, "CLIP")
                    view.models.append(ModelEntry(node=label, kind=kind, name=inputs[key]))

        if _LORA_RE.search(class_type):
            # simple case: lora_name + strength(s)
            if "lora_name" in inputs:
                strength = inputs.get("strength_model", inputs.get("strength", ""))
                view.models.append(ModelEntry(
                    node=label, kind="LoRA", name=inputs["lora_name"],
                    detail=f"weight {strength}" if strength != "" else "",
                ))
            # "Power Lora Loader" case / lists of {on, lora, strength} dicts
            for key, val in inputs.items():
                if isinstance(val, dict) and "lora" in val:
                    if val.get("on", True) is False:
                        continue
                    strength = val.get("strength", val.get("strength_model", ""))
                    view.models.append(ModelEntry(
                        node=label, kind="LoRA", name=val["lora"],
                        detail=f"weight {strength}" if strength != "" else "",
                    ))

    # --- Main parameters --------------------------------------------
    PARAM_KEYS = ["seed", "steps", "cfg", "sampler_name", "scheduler", "denoise",
                  "width", "height", "batch_size", "seconds", "duration",
                  "shift", "bpm", "timesignature", "language", "keyscale"]
    for node_id, node in prompt_graph.items():
        class_type = node.get("class_type", "")
        inputs = node.get("inputs", {})
        if not (_PARAM_CLASS_RE.search(class_type) or _ENCODE_CLASS_RE.search(class_type)):
            continue
        label = _node_label(node_id, workflow_by_id, class_type)
        for key in PARAM_KEYS:
            if key in inputs and not isinstance(inputs[key], (list, dict)):
                view.params.append(ParamEntry(node=label, param=key, value=str(inputs[key])))

    # --- Prompts (final resolved + dynamic template) ---------------------
    for node_id, node in prompt_graph.items():
        class_type = node.get("class_type", "")
        if not _ENCODE_CLASS_RE.search(class_type):
            continue
        inputs = node.get("inputs", {})
        label = _node_label(node_id, workflow_by_id, class_type)
        for key, role_label in _ROLE_LABELS.items():
            if key not in inputs:
                continue
            val = inputs[key]
            template = None
            source_node = None
            if isinstance(val, str):
                resolved = val
            elif isinstance(val, list):
                resolved, source_node = _resolve_link_text(prompt_graph, workflow_by_id, val)
                resolved = resolved or "(value computed at generation time, not stored in the file)"
            else:
                continue

            # Is the source node a (or does it contain a) dynamic prompt / wildcard?
            check_id = source_node or node_id
            src_prompt_node = prompt_graph.get(check_id, {})
            src_class = src_prompt_node.get("class_type", "")
            raw_widget = _first_string_widget(workflow_by_id.get(check_id))
            if (_DYNAMIC_CLASS_RE.search(src_class) or (raw_widget and _BRACE_PATTERN_RE.search(raw_widget))):
                if raw_widget and _BRACE_PATTERN_RE.search(raw_widget):
                    template = raw_widget

            view.prompts.append(PromptEntry(
                node=label, role=role_label, resolved=resolved,
                template=template, source_node=source_node,
            ))

    # --- "Orphan" dynamic prompt nodes (detected but not attached) --
    linked_sources = {p.source_node for p in view.prompts if p.source_node}
    for node_id, node in prompt_graph.items():
        class_type = node.get("class_type", "")
        if node_id in linked_sources:
            continue
        if not _DYNAMIC_CLASS_RE.search(class_type):
            raw_widget = _first_string_widget(workflow_by_id.get(node_id))
            if not (raw_widget and _BRACE_PATTERN_RE.search(raw_widget)):
                continue
        label = _node_label(node_id, workflow_by_id, class_type)
        template = _first_string_widget(workflow_by_id.get(node_id))
        inputs = node.get("inputs", {})
        resolved = None
        for key in _TEXT_KEY_PRIORITY:
            if key in inputs and isinstance(inputs[key], str):
                resolved = inputs[key]
                break
        if template and _BRACE_PATTERN_RE.search(template):
            view.prompts.append(PromptEntry(
                node=label, role="Dynamic prompt (unlinked)",
                resolved=resolved or "(not found)", template=template,
            ))

    # --- Deduplicate ------------------------------------------------
    # ComfyUI graphs commonly feed the exact same text into two parallel
    # branches (e.g. a positive pass and a separate "zeroed" / negative
    # pass both reading the same tags/lyrics node). That produces two
    # PromptEntry objects with identical content under the same role —
    # collapse those so a file with several *distinct* dynamic prompts
    # (e.g. one for music tags, one for lyrics, in Ace Step) shows one
    # card per distinct prompt instead of duplicates.
    seen: set[tuple[str, str, str | None]] = set()
    deduped: list[PromptEntry] = []
    for p in view.prompts:
        signature = (p.role, p.resolved, p.template)
        if signature in seen:
            continue
        seen.add(signature)
        deduped.append(p)
    view.prompts = deduped

    return view


def analyze(path: str | Path) -> StructuredView:
    raw = extract_raw_metadata(path)
    return build_structured_view(raw)


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("Usage: python3 comfyui_meta.py <file>")
        raise SystemExit(1)

    view = analyze(sys.argv[1])

    print("\n=== MODEL & LORAS ===")
    for m in view.models:
        print(f"[{m.kind}] {m.name}  ({m.node}) {m.detail}")

    print("\n=== MAIN PARAMETERS ===")
    for p in view.params:
        print(f"{p.node} :: {p.param} = {p.value}")

    print("\n=== PROMPT & DYNAMIC PROMPT ===")
    for pr in view.prompts:
        print(f"\n-- {pr.node} :: {pr.role} --")
        print("Resolved:", pr.resolved[:300])
        if pr.template:
            print("Dynamic template:", pr.template[:300])

    if view.warnings:
        print("\n=== WARNINGS ===")
        for w in view.warnings:
            print("-", w)