"""comfyui_score_clip_alignment — fidelidad prompt<->imagen via CLIP (0-1).

Juez de fidelidad del panel `comfyui-judge`: mide cuanto se parece el contenido de
una imagen al prompt que se pidio, usando el espacio compartido de CLIP. Embebe la
imagen y el texto con CLIP ViT-L/14, los normaliza y devuelve la similitud coseno
(clamp a 0-1). Valores tipicos: ~0.28-0.35 buen match, ~0.10-0.18 contenido distinto.

Impura: la inferencia necesita torch + open_clip, que NO estan en el venv del registry
(`python/.venv`). Por eso delega a un SUBPROCESO contra el python del venv de ComfyUI
(`~/ComfyUI/.venv/bin/python3`). Se calcula localmente con open_clip (no via el server
ComfyUI): es el camino mas simple y fiable, reusa el mismo CLIP que el scorer estetico
y no depende del estado del grafo del server. El parametro `server` queda reservado para
un backend futuro basado en los nodos CLIPVisionEncode/CLIPTextEncode (ver Gotchas).
"""
import json
import os
import subprocess
import sys

DEFAULT_VENV_PYTHON = "~/ComfyUI/.venv/bin/python3"

_MARKER = "__CLIP_RESULT__"
_INFER_SCRIPT = r'''
import json, sys
import torch
import open_clip
from PIL import Image

MARKER = "__CLIP_RESULT__"


def run(image_path, prompt, device):
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    model, _, preprocess = open_clip.create_model_and_transforms(
        "ViT-L-14-quickgelu", pretrained="openai")
    model = model.to(device).eval()
    tokenizer = open_clip.get_tokenizer("ViT-L-14-quickgelu")

    img = preprocess(Image.open(image_path).convert("RGB")).unsqueeze(0).to(device)
    txt = tokenizer([prompt]).to(device)
    with torch.no_grad():
        imf = model.encode_image(img)
        tf = model.encode_text(txt)
        imf = imf / imf.norm(dim=-1, keepdim=True)
        tf = tf / tf.norm(dim=-1, keepdim=True)
        cos = float((imf @ tf.T).cpu().numpy()[0][0])
    return max(0.0, min(1.0, cos))


if __name__ == "__main__":
    image_path, prompt, device = sys.argv[1], sys.argv[2], sys.argv[3]
    try:
        score = run(image_path, prompt, device)
        print(MARKER + json.dumps({"ok": True, "score_0_1": score, "error": ""}))
    except Exception as exc:
        print(MARKER + json.dumps({"ok": False, "score_0_1": 0.0, "error": repr(exc)}))
'''


def comfyui_score_clip_alignment(
    image_path: str,
    prompt: str,
    *,
    server: str = "127.0.0.1:8188",
    venv_python: str = DEFAULT_VENV_PYTHON,
    device: str = "auto",
    timeout: float = 180.0,
) -> dict:
    """Mide la fidelidad prompt<->imagen como similitud coseno CLIP (0-1).

    Args:
        image_path: ruta a la imagen en disco local.
        prompt: texto contra el que medir la alineacion (idealmente en ingles, como CLIP).
        server: host:port del server ComfyUI. RESERVADO para un backend futuro; hoy el
            calculo es local via open_clip (ver Gotchas). keyword-only.
        venv_python: python del venv de ComfyUI con torch + open_clip; se expande `~`.
            keyword-only.
        device: "auto" (cuda si hay, si no cpu), "cuda" o "cpu". keyword-only.
        timeout: timeout del subproceso en segundos. keyword-only.

    Returns:
        dict ``{ok, score_0_1, error}``. En exito ``ok=True`` y ``score_0_1`` es la
        similitud coseno clamp a [0,1] (mayor = mas fiel al prompt). En error ``ok=False``,
        ``score_0_1=0.0`` y ``error`` describe la causa. Nunca lanza excepcion.
    """
    img = os.path.expanduser(image_path)
    if not os.path.isfile(img):
        return {"ok": False, "score_0_1": 0.0, "error": f"imagen no encontrada: {img}"}
    if not prompt or not prompt.strip():
        return {"ok": False, "score_0_1": 0.0, "error": "prompt vacio"}

    py = os.path.expanduser(venv_python)
    if not os.path.isfile(py):
        return {"ok": False, "score_0_1": 0.0,
                "error": f"python del venv ComfyUI no encontrado: {py}"}

    try:
        proc = subprocess.run(
            [py, "-c", _INFER_SCRIPT, img, prompt, device],
            capture_output=True, text=True, timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        return {"ok": False, "score_0_1": 0.0,
                "error": f"subproceso de inferencia excedio {timeout}s"}

    for line in proc.stdout.splitlines():
        if line.startswith(_MARKER):
            try:
                return json.loads(line[len(_MARKER):])
            except json.JSONDecodeError as exc:
                return {"ok": False, "score_0_1": 0.0,
                        "error": f"JSON invalido del subproceso: {exc}"}

    tail = (proc.stderr or proc.stdout or "")[-400:]
    return {"ok": False, "score_0_1": 0.0,
            "error": f"sin resultado del subproceso (rc={proc.returncode}): {tail}"}


def _main(argv):
    if len(argv) < 2:
        sys.stderr.write('uso: comfyui_score_clip_alignment IMAGEN "prompt" [--device cpu|cuda]\n')
        return 2
    device = "auto"
    image_path = argv[0]
    prompt_parts = []
    i = 1
    while i < len(argv):
        if argv[i] == "--device" and i + 1 < len(argv):
            device = argv[i + 1]
            i += 2
        else:
            prompt_parts.append(argv[i])
            i += 1
    res = comfyui_score_clip_alignment(image_path, " ".join(prompt_parts), device=device)
    print(json.dumps(res, indent=2, ensure_ascii=False))
    return 0 if res["ok"] else 1


if __name__ == "__main__":
    sys.exit(_main(sys.argv[1:]))
