#!/usr/bin/env python3
"""One-shot Qwen-Image-2.1 evaluation with guaranteed profile restoration."""

from __future__ import annotations

import argparse
import json
import os
import pathlib
import random
import time
import urllib.parse
import urllib.request


CONTROLLER = os.environ.get("CONTROLLER_URL", "http://profile-controller:8090")
TOKEN = os.environ["CONTROLLER_TOKEN"]
COMFY = os.environ.get("COMFY_URL", "http://qwen-image-21-test:8188")
OUTPUT = pathlib.Path(os.environ.get("OUTPUT_DIR", "/output"))


def request_json(url: str, *, payload: dict | None = None,
                 authenticated: bool = False, timeout: float = 30) -> dict:
    body = None if payload is None else json.dumps(payload).encode()
    headers = {"Content-Type": "application/json"}
    if authenticated:
        headers["Authorization"] = f"Bearer {TOKEN}"
    method = "POST" if payload is not None else "GET"
    req = urllib.request.Request(url, data=body, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as response:
        return json.load(response)


def controller(path: str, *, post: bool = False) -> dict:
    return request_json(CONTROLLER + path, payload={} if post else None,
                        authenticated=True, timeout=900)


def wait_comfy(timeout: int = 900) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            request_json(COMFY + "/system_stats", timeout=3)
            return
        except Exception:
            time.sleep(2)
    raise TimeoutError("Qwen-Image worker did not become ready")


def workflow(prompt: str, seed: int, steps: int, width: int, height: int) -> dict:
    return {
        "1": {"class_type": "UNETLoader", "inputs": {
            "unet_name": "qwen_image_2.1_int8_convrot.safetensors",
            "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "qwen3vl_8b_int8_convrot.safetensors",
            "type": "qwen_image", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {
            "vae_name": "qwen_image_2.1_vae_bf16.safetensors"}},
        "4": {"class_type": "TextEncodeQwenImage21", "inputs": {
            "clip": ["2", 0], "prompt": prompt, "negative_prompt": "",
            "resolution": max(width, height)}},
        "5": {"class_type": "EmptyLatentImage", "inputs": {
            "width": width, "height": height, "batch_size": 1}},
        "6": {"class_type": "KSampler", "inputs": {
            "model": ["1", 0], "positive": ["4", 0], "negative": ["4", 1],
            "latent_image": ["5", 0], "seed": seed, "steps": steps,
            "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple",
            "denoise": 1.0}},
        "7": {"class_type": "VAEDecode", "inputs": {
            "samples": ["6", 0], "vae": ["3", 0]}},
        "8": {"class_type": "SaveImage", "inputs": {
            "filename_prefix": "qwen-image-2.1-test", "images": ["7", 0]}},
    }


def wait_result(prompt_id: str, timeout: int = 3600) -> dict:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        history = request_json(f"{COMFY}/history/{prompt_id}", timeout=10)
        if prompt_id in history:
            result = history[prompt_id]
            status = result.get("status", {})
            if status.get("status_str") == "error" or not status.get("completed", True):
                raise RuntimeError("generation failed: " + json.dumps(status))
            return result
        time.sleep(2)
    raise TimeoutError("Qwen-Image generation timed out")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("prompt")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--steps", type=int, default=25)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--height", type=int, default=1024)
    args = parser.parse_args()
    if args.width % 32 or args.height % 32:
        parser.error("width and height must be multiples of 32")
    seed = args.seed if args.seed is not None else random.randrange(2**53)
    previous = controller("/profiles/status").get("active_profile")
    started = time.monotonic()
    try:
        controller("/workers/qwen-image-test/start", post=True)
        wait_comfy()
        queued = request_json(COMFY + "/prompt", payload={
            "prompt": workflow(args.prompt, seed, args.steps, args.width, args.height),
            "client_id": "athena-qwen-image-21-test"}, timeout=30)
        prompt_id = queued["prompt_id"]
        result = wait_result(prompt_id)
        images = []
        for node in result.get("outputs", {}).values():
            images.extend(node.get("images", []))
        if not images:
            raise RuntimeError("generation completed without an image")
        image = images[0]
        query = urllib.parse.urlencode({
            "filename": image["filename"], "subfolder": image.get("subfolder", ""),
            "type": image.get("type", "output")})
        target = OUTPUT / image["filename"]
        target.parent.mkdir(parents=True, exist_ok=True)
        with urllib.request.urlopen(COMFY + "/view?" + query, timeout=120) as src:
            target.write_bytes(src.read())
        print(json.dumps({"status": "ok", "file": str(target), "seed": seed,
                          "seconds": round(time.monotonic() - started, 1)}, indent=2))
    finally:
        try:
            controller("/workers/qwen-image-test/stop", post=True)
        finally:
            if previous:
                controller(f"/profiles/{previous}/activate", post=True)


if __name__ == "__main__":
    main()