ComfyAssets-ComfyUI_PromptManager/utils/prompt_inspector.py

"""Read prompt evidence without inference, indexing, or decoding image pixels."""

import hashlib
import io
import json
import re
import struct
import subprocess  # nosec B404
import threading
import zlib
from dataclasses import dataclass, field as dataclass_field
from pathlib import Path
from typing import Any

from PIL import Image

from .inspection_outcomes import UNREADABLE_METADATA_REASON, outcome_message
from .pair_execution import EXECUTION_KEY
from .prompt_pair_resolver import (
    _FOOTER,
    EMPTY_POSITIVE_REASON,
    INSPECTION_PROFILE,
    ResolverProfile,
    resolve_prompt_pair,
)
from .video_metadata import ffprobe_path, is_video_path

# Role records (prompt, parameters, the pair execution record) share one
# budget, encoded and decoded: roles cannot be trusted from a partial record,
# so exceeding it fails the whole read. A workflow has its own budget; past it
# only the workflow is ``limited``. Other text chunks are never retained.
ROLE_METADATA_BYTES = 2 * 1024 * 1024
WORKFLOW_METADATA_BYTES = 4 * 1024 * 1024
# Total bytes read for byte-identical duplicate chunks of one file, per class
# so workflow repeats can never starve role records: two budgets each, making
# repeated chunks cost bounded I/O instead of file size.
DUPLICATE_WORKFLOW_BYTES = 2 * WORKFLOW_METADATA_BYTES
DUPLICATE_ROLE_BYTES = 2 * ROLE_METADATA_BYTES
MAX_METADATA_BYTES = ROLE_METADATA_BYTES  # video stdout bound, unchanged here
MAX_KEYWORD_BYTES = 80  # a 1-79 byte PNG keyword plus its NUL separator
MAX_HEADER_BYTES = 8 * 1024 * 1024
MAX_PNG_CHUNKS = 10_000
_ROLE_METADATA_KEYS = frozenset({"prompt", "parameters", EXECUTION_KEY})
_METADATA_KEYS = _ROLE_METADATA_KEYS | {"workflow"}
_TEXT_CHUNKS = (b"tEXt", b"zTXt", b"iTXt")
# An accumulating batch saver writes one A1111 ``parameters`` record per batch
# item, identical except this field (R3); Pillow keeps the last one.
# Matched only inside the record's one ``Steps: N, ...`` settings footer line.
_BATCH_INDEX = re.compile(r"(?:^|(?<=, ))Batch index: \d+(?=, |$)")


def _batch_normalized(text: str) -> str | None:
    """``text`` with its footer's single batch index blanked, else ``None``."""
    footers = list(_FOOTER.finditer(text))
    if len(footers) != 1:
        return None
    line, count = _BATCH_INDEX.subn("Batch index: #", footers[0].group())
    if count != 1:
        return None
    return text[: footers[0].start()] + line + text[footers[0].end() :]


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


class MetadataReadError(ValueError):
    """Prompt evidence is unreadable, incomplete, or exceeds inspection limits."""


@dataclass(frozen=True)
class InspectionRead:
    """One bounded read, with its transient exact video source kept private."""

    metadata: dict[str, str]
    record_formats: dict[str, str]
    video_carrier_text: str | None
    format: str | None
    width: int | None
    height: int | None
    # Records present but over their budget: no text, format ``limited``.
    limited_records: frozenset[str] = dataclass_field(default_factory=frozenset)


@dataclass(frozen=True)
class _VideoInspection:
    metadata: dict[str, str | None]
    record_formats: dict[str, str]
    carrier_text: str | None


def _inflate(data: bytes, limit: int) -> bytes | None:
    """Inflate at most ``limit + 1`` bytes; ``None`` when the output is over limit."""
    inflater = zlib.decompressobj()
    result = inflater.decompress(data, limit + 1)
    if len(result) > limit:
        return None
    if not inflater.eof or inflater.unused_data:
        raise MetadataReadError("Compressed metadata is invalid or exceeds limits")
    return result


def _keyword(prefix: bytes) -> str:
    keyword, separator, _ = prefix.partition(b"\0")
    if not separator or not 1 <= len(keyword) <= 79:
        raise MetadataReadError("Malformed PNG text metadata")
    return keyword.decode("latin-1")


def _png_text(kind: bytes, data: bytes, limit: int) -> tuple[str, str | None]:
    """Decode one retained text chunk; ``None`` text when it inflates past limit."""
    _, _, body = data.partition(b"\0")
    key = _keyword(data[:MAX_KEYWORD_BYTES])
    encoding = "latin-1"
    if kind == b"zTXt":
        if not body or body[0] != 0:
            raise MetadataReadError("Unsupported PNG text compression")
        body = _inflate(body[1:], limit)
    elif kind == b"iTXt":
        if len(body) < 2 or body[0] not in (0, 1):
            raise MetadataReadError("Malformed international PNG metadata")
        flag, method = body[:2]
        fields = body[2:].split(b"\0", 2)
        if len(fields) != 3 or (flag and method != 0):
            raise MetadataReadError("Malformed international PNG metadata")
        fields[0].decode("ascii")
        fields[1].decode("utf-8")
        body = _inflate(fields[2], limit) if flag else fields[2]
        encoding = "utf-8"
    if body is None:
        return key, None
    if b"\0" in body:
        raise MetadataReadError("Malformed PNG text metadata")
    return key, body.decode(encoding)


class _PngRecords:
    """Retained records with per-class budgets and byte-identical duplicates."""

    def __init__(self) -> None:
        self.found: dict[str, str] = {}
        self.limited: set[str] = set()
        self.digests: dict[str, tuple[int, bytes]] = {}
        self.role_encoded = self.role_decoded = 0
        self.duplicate_workflow = self.duplicate_role = 0
        self.duplicate_decoded = 0

    def skip(self, key: str, length: int) -> bool:
        """Decide from the keyword and length alone whether to seek past a chunk."""
        if key not in _METADATA_KEYS or key in self.limited:
            return True  # never retained, or a duplicate of a limited workflow
        if key in self.digests:
            if self.digests[key][0] != length and key != "parameters":
                raise MetadataReadError("Duplicate prompt metadata")
            if key == "workflow":
                if self.duplicate_workflow + length > DUPLICATE_WORKFLOW_BYTES:
                    self.limited.add(key)
                    self.found.pop(key, None)
                    return True
                self.duplicate_workflow += length
            else:
                if self.duplicate_role + length > DUPLICATE_ROLE_BYTES:
                    raise MetadataReadError("Metadata exceeds inspection limits")
                self.duplicate_role += length
            return False  # same length: read and compare bytes
        if key == "workflow":
            if length > WORKFLOW_METADATA_BYTES:
                self.limited.add(key)
                return True
            return False
        self.role_encoded += length
        if self.role_encoded > ROLE_METADATA_BYTES:
            raise MetadataReadError("Metadata exceeds inspection limits")
        return False

    def add(self, kind: bytes, data: bytes) -> None:
        key = _keyword(data[:MAX_KEYWORD_BYTES])
        digest = hashlib.sha256(kind + data).digest()
        if key in self.digests:
            if self.digests[key][1] != digest:
                if key != "parameters":
                    raise MetadataReadError("Duplicate prompt metadata")
                self.batch_copy(kind, data)
                # The kept copy is the last one; later repeats compare with it.
                self.digests[key] = (len(data), digest)
            return
        self.digests[key] = (len(data), digest)
        workflow = key == "workflow"
        limit = (
            WORKFLOW_METADATA_BYTES
            if workflow
            else ROLE_METADATA_BYTES - self.role_decoded
        )
        _, value = _png_text(kind, data, limit)
        size = None if value is None else len(value.encode("utf-8"))
        if workflow:
            if size is None or size > WORKFLOW_METADATA_BYTES:
                self.limited.add(key)
                return
        else:
            self.role_decoded += limit + 1 if size is None else size
            if self.role_decoded > ROLE_METADATA_BYTES:
                raise MetadataReadError("Metadata exceeds inspection limits")
        self.found[key] = value

    def batch_copy(self, kind: bytes, data: bytes) -> None:
        """Keep a later ``parameters`` copy that differs only in its batch index.

        Each copy must carry exactly one ``Batch index: N`` settings field; any
        other difference is still a conflict. Decoding every copy draws from the
        duplicate role budget, so small compressed copies stay bounded.
        """
        _, value = _png_text(kind, data, ROLE_METADATA_BYTES)
        size = ROLE_METADATA_BYTES + 1 if value is None else len(value.encode())
        self.duplicate_decoded += size
        if value is None or self.duplicate_decoded > DUPLICATE_ROLE_BYTES:
            raise MetadataReadError("Metadata exceeds inspection limits")
        kept = self.found["parameters"]
        normalized = _batch_normalized(kept)
        if normalized is None or normalized != _batch_normalized(value):
            raise MetadataReadError("Duplicate prompt metadata")
        self.role_decoded += max(0, size - len(kept.encode()))
        if self.role_decoded > ROLE_METADATA_BYTES:
            raise MetadataReadError("Metadata exceeds inspection limits")
        self.found["parameters"] = value


def _png_metadata(stream: Any, size: int) -> tuple[dict, frozenset, int, int]:
    """CRC-check retained text and structural chunks, seeking past everything else.

    PNG text chunks may follow IDAT. Pillow's ``.text`` loads pixels to find
    them; this bounded reader instead follows the W3C PNG chunk layout. A text
    chunk's keyword is read first (at most ``MAX_KEYWORD_BYTES``); chunks that
    are not retained, or are over their budget, are skipped by length, so their
    bytes are never held in memory or checksummed. IDAT checksums/pixel
    validity are deliberately outside metadata inspection.
    """
    records = _PngRecords()
    saw_header = saw_pixels = False
    for index in range(MAX_PNG_CHUNKS):
        header = stream.read(8)
        if len(header) != 8:
            raise MetadataReadError("Incomplete PNG metadata")
        length, kind = struct.unpack(">I4s", header)
        if length > 0x7FFFFFFF or stream.tell() + length + 4 > size:
            raise MetadataReadError("Incomplete PNG chunk")
        if index == 0 and (kind != b"IHDR" or length != 13):
            raise MetadataReadError("Missing PNG header")
        if kind in _TEXT_CHUNKS:
            prefix = stream.read(min(length, MAX_KEYWORD_BYTES))
            if records.skip(_keyword(prefix), length):
                stream.seek(length - len(prefix) + 4, 1)
                continue
            data = prefix + stream.read(length - len(prefix))
        elif kind in (b"IHDR", b"IEND"):
            data = stream.read(length)
        else:
            saw_pixels |= kind == b"IDAT"
            stream.seek(length + 4, 1)
            continue
        checksum = stream.read(4)
        if zlib.crc32(kind + data) != int.from_bytes(checksum, "big"):
            raise MetadataReadError("PNG metadata checksum mismatch")
        if kind == b"IHDR":
            if saw_header or not all(struct.unpack(">II", data[:8])):
                raise MetadataReadError("Malformed PNG header")
            saw_header = True
            width, height = struct.unpack(">II", data[:8])
        elif kind == b"IEND":
            if length or not saw_pixels:
                raise MetadataReadError("Incomplete PNG image")
            return records.found, frozenset(records.limited), width, height
        else:
            records.add(kind, data)
    raise MetadataReadError("PNG chunk count exceeds inspection limits")


def _video_metadata(path: Path) -> _VideoInspection:
    probe = ffprobe_path()
    if not probe:
        raise MetadataReadError("Video metadata reader is unavailable")
    command = [
        probe,
        "-v",
        "quiet",
        "-show_entries",
        "format_tags=comment,description",
        "-of",
        "json",
        str(path),
    ]
    process = subprocess.Popen(  # nosec B603
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
    )

    def stop() -> None:
        if process.poll() is None:
            try:
                process.kill()
            except OSError:
                # Timeout and final cleanup can race with the process exiting.
                pass

    timer = threading.Timer(3, stop)
    timer.daemon = True
    timer.start()
    try:
        raw = process.stdout.read(MAX_METADATA_BYTES + 1)
        if len(raw) > MAX_METADATA_BYTES:
            raise MetadataReadError("Video metadata exceeds inspection limits")
        process.wait()
        if process.returncode:
            raise MetadataReadError("Video metadata could not be read")
    finally:
        timer.cancel()
        stop()
        process.wait()
        process.stdout.close()
    probed = json.loads(raw)
    if not isinstance(probed, dict) or not isinstance(probed.get("format"), dict):
        raise MetadataReadError("Malformed video metadata")
    tags = probed["format"].get("tags", {})
    if not isinstance(tags, dict):
        raise MetadataReadError("Malformed video metadata")
    # The scan normalizer drops null fields and unknown keys. Inspection must
    # retain explicit prompt/execution markers so malformed evidence cannot
    # silently fall back to an unrelated saved recipe or another comment.
    candidates = []
    for key in ("comment", "description"):
        raw = tags.get(key)
        if not isinstance(raw, str):
            continue
        try:
            envelope = json.loads(raw)
        except ValueError:
            continue
        if not isinstance(envelope, dict):
            continue
        candidates.append((raw, envelope))
    selected = next(
        (item for item in candidates if _ROLE_METADATA_KEYS.intersection(item[1])),
        None,
    )
    if selected is None:
        selected = next((item for item in candidates if "workflow" in item[1]), None)
    if selected is not None:
        raw, envelope = selected
        metadata, formats = {}, {}
        for field in _METADATA_KEYS.intersection(envelope):
            value = envelope[field]
            if field == "workflow" and not isinstance(value, (str, dict)):
                formats[field] = "invalid"
                continue
            metadata[field] = (
                value if value is None or isinstance(value, str) else json.dumps(value)
            )
            formats[field] = (
                "original-text" if isinstance(value, str) else "reserialized-json"
            )
        return _VideoInspection(metadata, formats, raw)
    return _VideoInspection({}, {}, None)


def read_inspection_file(path: Path) -> InspectionRead:
    """Read bounded metadata and header facts without decoding media pixels."""
    try:
        formats, carrier, limited = {}, None, frozenset()
        format_name = width = height = None
        if is_video_path(path):
            video = _video_metadata(path)
            metadata, formats, carrier = (
                video.metadata,
                video.record_formats,
                video.carrier_text,
            )
        else:
            with path.open("rb") as stream:
                signature = stream.read(8)
                if signature == _SIGNATURE:
                    metadata, limited, width, height = _png_metadata(
                        stream, path.stat().st_size
                    )
                    format_name = "PNG"
                else:
                    stream.seek(0)
                    header = stream.read(MAX_HEADER_BYTES)
                    with Image.open(io.BytesIO(header)) as image:
                        metadata = dict(image.info)
                        format_name, (width, height) = image.format, image.size
        result, role_size, limited = {}, 0, set(limited)
        for key in _METADATA_KEYS:
            if key not in metadata:
                continue
            value = metadata[key]
            if key == "workflow" and not isinstance(value, str):
                formats[key] = "invalid"
                continue
            if not isinstance(value, str):
                raise MetadataReadError("Malformed prompt metadata")
            size = len(value.encode("utf-8", "surrogatepass"))
            if key == "workflow":
                if size > WORKFLOW_METADATA_BYTES:
                    limited.add(key)
                    continue
            else:
                role_size += size
                if role_size > ROLE_METADATA_BYTES:
                    raise MetadataReadError("Metadata exceeds inspection limits")
            result[key] = value
            formats.setdefault(key, "original-text")
        if _ROLE_METADATA_KEYS.intersection(limited):
            # A limited role record would vanish and let fallback evidence win.
            raise MetadataReadError("Metadata exceeds inspection limits")
        for key in limited:
            formats[key] = "limited"
        return InspectionRead(
            result,
            formats,
            carrier,
            format_name,
            width,
            height,
            frozenset(limited),
        )
    except MetadataReadError:
        raise
    except (
        OSError,
        ValueError,
        EOFError,
        SyntaxError,
        RecursionError,
        struct.error,
        zlib.error,
    ):
        raise MetadataReadError("Media metadata could not be read") from None


def read_inspection_metadata(path: Path) -> dict[str, str]:
    """Compatible metadata-only projection of the single inspection read."""
    return read_inspection_file(path).metadata


def _side(status: str, text: str | None = None, source: str | None = None) -> dict:
    return {"status": status, "text": text, "source": source}


_EMPTY_POSITIVE = outcome_message(
    {"state": "unresolved", "reason": EMPTY_POSITIVE_REASON}
)


def _explained(prompts: dict, resolution: dict) -> dict:
    """Each unresolved role carries the outcome's line; other roles are kept."""
    message = outcome_message(resolution)
    if message is None:
        return prompts
    return {
        role: (
            {**side, "message": message}
            if side["status"] == "unresolved" and "message" not in side
            else side
        )
        for role, side in prompts.items()
    }


def unreadable_evidence() -> dict:
    """Both roles when the file's metadata could not be read at all."""
    resolution = {
        "state": "unresolved",
        "sampler_id": None,
        "reason": UNREADABLE_METADATA_REASON,
        "reason_code": None,
        "node": None,
    }
    prompts = {role: _side("unresolved") for role in ("positive", "negative")}
    return {"prompts": _explained(prompts, resolution), "resolution": resolution}


def _may_hold_prompt_manager(metadata: dict) -> bool:
    """Whether the API graph can show the R1 signature, so the workflow is needed.

    PromptManager classes, or a ``CLIPTextEncode`` carrying PromptManager's own
    ``search_text`` input (an older save rewrote the class).
    """
    graph = metadata.get("prompt")
    if not isinstance(graph, dict):
        return False
    for node in graph.values():
        if not isinstance(node, dict):
            continue
        name = node.get("class_type")
        inputs = node.get("inputs")
        if isinstance(name, str) and "promptmanager" in name.lower():
            return True
        if isinstance(inputs, dict) and "search_text" in inputs:
            return True
    return False


def _parsed_workflow(value: Any) -> dict | None:
    """The UI workflow as an object, or ``None`` (absent, invalid or not JSON)."""
    if isinstance(value, str):
        try:
            value = json.loads(value)
        except (ValueError, RecursionError):
            return None
    return value if isinstance(value, dict) else None


def inspect_prompt_evidence(
    metadata: dict,
    saved: dict | None = None,
    *,
    profile: ResolverProfile = INSPECTION_PROFILE,
) -> dict:
    """Keep uncertainty visible; a saved association never replaces file evidence.

    Inspection routes resolve with ``INSPECTION_PROFILE``: bounded work, not
    graph size. Ingestion (scan, live monitor, admin import) calls the resolver
    directly with its default profile. ``resolution`` carries the closed
    ``reason_code`` and, for a blocking node, ``node: {id, class_type}``.
    """
    workflow = metadata.get("workflow")
    metadata = {
        key: value for key, value in metadata.items() if key in _ROLE_METADATA_KEYS
    }
    resolution = {
        "state": "unresolved",
        "sampler_id": None,
        "reason": "Prompt evidence is malformed",
        "reason_code": "malformed",
        "node": None,
    }
    if not metadata:
        result = {role: _side("absent") for role in ("positive", "negative")}
        if saved:
            context = saved.get("pair_context")
            paired = isinstance(saved.get("pair"), dict) or isinstance(context, dict)
            for role in ("positive", "negative") if paired else ("positive",):
                result[role] = {
                    **_side("unresolved", source="saved-association"),
                    "message": (
                        "A saved library association exists, but it cannot be "
                        "verified against this file."
                    ),
                }
        return {
            "prompts": result,
            "resolution": {
                "state": "absent",
                "sampler_id": None,
                "reason": "No parsed prompt metadata",
                "reason_code": None,
                "node": None,
            },
        }
    canonical = dict(metadata)
    malformed = {
        "prompts": _explained(
            {role: _side("unresolved") for role in ("positive", "negative")},
            resolution,
        ),
        "resolution": resolution,
    }
    if "prompt" not in canonical and "parameters" not in canonical:
        return malformed
    if isinstance(canonical.get("prompt"), str):
        try:
            canonical["prompt"] = json.loads(canonical["prompt"])
        except (ValueError, RecursionError):
            return malformed
    workflow = (
        _parsed_workflow(workflow) if _may_hold_prompt_manager(canonical) else None
    )
    if workflow is not None and "prompt" in canonical:
        # Only to recognize overwritten PromptManager text (R1); never a role.
        canonical["workflow"] = workflow
    resolved = resolve_prompt_pair(canonical, profile=profile)
    state = resolved["state"]
    result = {}
    for role in ("positive", "negative"):
        side = resolved[role]
        if (
            state != "ambiguous"
            and isinstance(side, dict)
            and (role == "negative" or side["text"].strip())
        ):
            result[role] = _side(
                "resolved", side["text"], side["provenance"]["text_source"]
            )
        elif state == "resolved" and isinstance(side, dict):
            # A resolved pair withholds only an empty positive.
            result[role] = {**_side("unresolved"), "message": _EMPTY_POSITIVE}
        else:
            result[role] = _side(state if state != "resolved" else "unresolved")
    resolution = {
        "state": state,
        "sampler_id": resolved.get("sampler_id"),
        "reason": resolved.get("reason", ""),
        "reason_code": resolved.get("reason_code"),
        "node": resolved.get("node"),
    }
    return {"prompts": _explained(result, resolution), "resolution": resolution}


def inspect_prompt_roles(metadata: dict, saved: dict | None = None) -> dict:
    """Keep the established role-only helper contract."""
    return inspect_prompt_evidence(metadata, saved)["prompts"]