"""MCP ComfyUI pour Hermes avec transmission d'images vers l'interface Alfred.

Fonctionne avec un workflow API ComfyUI possédant un nœud de sauvegarde d'image.
L'envoi HTTP à Alfred est indépendant de la réponse textuelle d'Hermes.
"""

import json
import os
import time
import uuid
from pathlib import Path
from urllib.parse import urlencode

import requests
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

load_dotenv(Path(__file__).with_name(".env"))

mcp = FastMCP("Alfred ComfyUI")
BASE = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
WF = Path(os.getenv("WORKFLOWS_DIR", str(Path(__file__).parent / "workflows")))
ALFRED_RENDER_URL = os.getenv(
    "ALFRED_RENDER_URL", "http://127.0.0.1:54321/api/render"
)
ALFRED_RENDER_TOKEN = os.getenv("ALFRED_RENDER_TOKEN", "").strip()


def submit(workflow: dict) -> str:
    response = requests.post(
        f"{BASE}/prompt",
        json={"prompt": workflow, "client_id": str(uuid.uuid4())},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()["prompt_id"]


def wait(prompt_id: str) -> dict:
    deadline = time.monotonic() + int(os.getenv("COMFYUI_TIMEOUT", "300"))
    while time.monotonic() < deadline:
        response = requests.get(f"{BASE}/history/{prompt_id}", timeout=15)
        response.raise_for_status()
        result = response.json()
        if prompt_id in result:
            return result[prompt_id]
        time.sleep(1)
    raise TimeoutError(f"ComfyUI timeout ({prompt_id})")


def extract_images(history: dict) -> list[dict]:
    """Construit les URL de /view depuis les sorties ComfyUI (SaveImage)."""
    images = []
    for node_id, output in history.get("outputs", {}).items():
        for image in output.get("images", []):
            filename = image.get("filename")
            if not filename:
                continue
            subfolder = image.get("subfolder", "")
            image_type = image.get("type", "output")
            query = urlencode(
                {
                    "filename": filename,
                    "subfolder": subfolder,
                    "type": image_type,
                }
            )
            images.append(
                {
                    "node_id": str(node_id),
                    "filename": filename,
                    "subfolder": subfolder,
                    "type": image_type,
                    "url": f"{BASE}/view?{query}",
                }
            )
    return images


def show_on_alfred(image: dict, caption: str) -> tuple[bool, str | None]:
    """Ne fait pas échouer la génération si Alfred est hors ligne."""
    headers = {}
    if ALFRED_RENDER_TOKEN:
        headers["X-Alfred-Render-Token"] = ALFRED_RENDER_TOKEN
    try:
        response = requests.post(
            ALFRED_RENDER_URL,
            json={
                "type": "image",
                "url": image["url"],
                "alt": "Image générée par ComfyUI",
                "caption": caption,
            },
            headers=headers,
            timeout=5,
        )
        response.raise_for_status()
        return True, None
    except requests.RequestException as exc:
        return False, str(exc)


@mcp.tool()
def generate_image(prompt: str) -> dict:
    """Génère une image par ComfyUI et l'affiche automatiquement sur Alfred.

    Utiliser cet outil lorsque l'utilisateur demande de créer/générer une image.
    Ne pas prétendre que l'image a été affichée si displayed_on_alfred est faux.
    """
    prompt = str(prompt).strip()
    if not prompt:
        raise ValueError("Le prompt de génération est vide")

    workflow = json.loads((WF / "text_to_image.json").read_text(encoding="utf-8"))
    prompt_node = os.getenv("COMFY_PROMPT_NODE", "6")
    workflow[prompt_node]["inputs"]["text"] = prompt

    prompt_id = submit(workflow)
    history = wait(prompt_id)
    images = extract_images(history)
    if not images:
        raise RuntimeError(
            "ComfyUI a terminé sans image de sortie. "
            "Vérifie le nœud SaveImage du workflow."
        )

    displayed, display_error = show_on_alfred(
        images[0], f"Image générée : {prompt}"
    )
    return {
        "ok": True,
        "prompt_id": prompt_id,
        "prompt": prompt,
        "images": images,
        "first_image_url": images[0]["url"],
        "displayed_on_alfred": displayed,
        "display_error": display_error,
    }


@mcp.tool()
def analyze_image(image_path: str, question: str) -> dict:
    """Analyse une image via le workflow VLM configuré (inchangé)."""
    workflow = json.loads((WF / "vision_analyze.json").read_text(encoding="utf-8"))
    workflow[os.getenv("COMFY_VISION_IMAGE_NODE", "1")]["inputs"]["image"] = image_path
    workflow[os.getenv("COMFY_VISION_PROMPT_NODE", "2")]["inputs"]["text"] = question
    return wait(submit(workflow))


if __name__ == "__main__":
    mcp.run()