import os, json, time, uuid, requests
from pathlib import Path
from urllib.parse import quote
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

load_dotenv()

mcp = FastMCP("Alfred ComfyUI")
BASE = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
WF = Path(os.getenv("WORKFLOWS_DIR", Path(__file__).parent / "workflows"))


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


def wait(pid):
    deadline = time.time() + int(os.getenv("COMFYUI_TIMEOUT", "300"))
    while time.time() < deadline:
        h = requests.get(f"{BASE}/history/{pid}", timeout=15).json()
        if pid in h:
            return h[pid]
        time.sleep(1)
    raise TimeoutError("ComfyUI timeout")


def build_image_url(filename, subfolder="", img_type="output"):
    return (
        f"{BASE}/view?"
        f"filename={quote(str(filename))}"
        f"&subfolder={quote(str(subfolder))}"
        f"&type={quote(str(img_type))}"
    )


def extract_images(history):
    images = []

    outputs = history.get("outputs", {})
    for node_id, node_output in outputs.items():
        for img in node_output.get("images", []):
            filename = img.get("filename")
            subfolder = img.get("subfolder", "")
            img_type = img.get("type", "output")

            if not filename:
                continue

            images.append({
                "node_id": node_id,
                "filename": filename,
                "subfolder": subfolder,
                "type": img_type,
                "url": build_image_url(filename, subfolder, img_type)
            })

    return images


@mcp.tool()
def generate_image(prompt: str):
    """Generate an image with the locked text_to_image workflow."""
    w = json.loads((WF / "text_to_image.json").read_text())
    node = os.getenv("COMFY_PROMPT_NODE", "6")
    w[node]["inputs"]["text"] = prompt

    history = wait(submit(w))
    images = extract_images(history)

    return {
        "ok": True,
        "message": f"Image générée pour : {prompt}",
        "prompt": prompt,
        "images": images,
        "first_image_url": images[0]["url"] if images else None,
        "first_image_filename": images[0]["filename"] if images else None
    }


@mcp.tool()
def analyze_image(image_path: str, question: str):
    """Analyze a camera/local image using the locked VLM workflow. Workflow node IDs are configured in .env."""
    w = json.loads((WF / "vision_analyze.json").read_text())
    w[os.getenv("COMFY_VISION_IMAGE_NODE", "1")]["inputs"]["image"] = image_path
    w[os.getenv("COMFY_VISION_PROMPT_NODE", "2")]["inputs"]["text"] = question
    return wait(submit(w))


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