"""Extrae el workflow embebido en los chunks de texto de un PNG de ComfyUI.

ComfyUI guarda en los PNG generados dos chunks de texto:
- "prompt": el workflow en API format (lo que se envio a POST /prompt).
- "workflow": el grafo de la UI (UI graph), presente si se genero desde la UI.

Lee chunks tEXt, zTXt e iTXt con stdlib (struct, zlib). Impura: red opcional (si
source es URL) + lectura de disco.
"""
import json
import struct
import urllib.error
import urllib.request
import zlib


def comfyui_import_workflow_png(png_path_or_url: str, timeout: float = 15.0) -> dict:
    """Devuelve el/los workflow(s) embebido(s) en un PNG de ComfyUI.

    Args:
        png_path_or_url: ruta local de un PNG, o URL http(s) de un PNG.
        timeout: timeout HTTP en segundos (solo si es URL).

    Returns:
        dict {ok, prompt, workflow, format_detected, error}:
          - prompt: API format (dict) si el chunk "prompt" existe, si no {}.
          - workflow: UI graph (dict) si el chunk "workflow" existe, si no {}.
          - format_detected: chunks hallados unidos por "+" ("prompt",
            "workflow" o "prompt+workflow").
        Si el PNG no trae metadata de workflow, ok=False.
    """
    try:
        if png_path_or_url.startswith(("http://", "https://")):
            with urllib.request.urlopen(png_path_or_url, timeout=timeout) as resp:
                data = resp.read()
        else:
            with open(png_path_or_url, "rb") as f:
                data = f.read()
    except (urllib.error.URLError, OSError) as exc:
        return {"ok": False, "prompt": {}, "workflow": {}, "format_detected": "",
                "error": f"no se pudo leer {png_path_or_url!r}: {exc}"}

    try:
        chunks = _png_text_chunks(data)
    except ValueError as exc:
        return {"ok": False, "prompt": {}, "workflow": {}, "format_detected": "",
                "error": str(exc)}

    out = {"ok": False, "prompt": {}, "workflow": {}, "format_detected": "", "error": ""}
    found = []
    if "prompt" in chunks:
        try:
            out["prompt"] = json.loads(chunks["prompt"])
            found.append("prompt")
        except json.JSONDecodeError:
            pass
    if "workflow" in chunks:
        try:
            out["workflow"] = json.loads(chunks["workflow"])
            found.append("workflow")
        except json.JSONDecodeError:
            pass
    out["format_detected"] = "+".join(found)
    if found:
        out["ok"] = True
    else:
        out["error"] = "el PNG no contiene metadata de workflow ComfyUI (chunks prompt/workflow)"
    return out


def _png_text_chunks(data: bytes) -> dict:
    """Lee los chunks de texto (tEXt/zTXt/iTXt) de un PNG -> {keyword: texto}."""
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError("no es un PNG valido (firma incorrecta)")
    out = {}
    off = 8
    n = len(data)
    while off + 8 <= n:
        length = struct.unpack(">I", data[off:off + 4])[0]
        ctype = data[off + 4:off + 8]
        body = data[off + 8:off + 8 + length]
        off += 12 + length  # 4 len + 4 type + body + 4 crc
        if ctype == b"tEXt":
            kw, _, txt = body.partition(b"\x00")
            out[kw.decode("latin1")] = txt.decode("latin1")
        elif ctype == b"zTXt":
            kw, _, rest = body.partition(b"\x00")
            if rest:
                comp_data = rest[1:]  # rest[0] = metodo de compresion
                try:
                    out[kw.decode("latin1")] = zlib.decompress(comp_data).decode("latin1")
                except zlib.error:
                    pass
        elif ctype == b"iTXt":
            kw, _, rest = body.partition(b"\x00")
            if len(rest) >= 2:
                comp_flag = rest[0]
                parts = rest[2:].split(b"\x00", 2)  # lang\x00 translated\x00 text
                if len(parts) == 3:
                    text_bytes = parts[2]
                    if comp_flag == 1:
                        try:
                            text_bytes = zlib.decompress(text_bytes)
                        except zlib.error:
                            text_bytes = b""
                    out[kw.decode("latin1")] = text_bytes.decode("utf-8", "replace")
        elif ctype == b"IEND":
            break
    return out


if __name__ == "__main__":
    import json as _json
    import sys

    path = sys.argv[1] if len(sys.argv) > 1 else "/tmp/missing.png"
    res = comfyui_import_workflow_png(path)
    print(_json.dumps({k: v for k, v in res.items() if k != "prompt"}, indent=2))
    print("nodos en prompt:", len(res["prompt"]))
