#!/usr/bin/env python3
"""
Quick CLI for generating images via ComfyUI + FLUX.2 Klein Heretic.

Usage:
    python gen.py "your prompt here"
    python gen.py "your prompt" --steps 20 --width 1280 --height 720
    python gen.py "your prompt" --seed 12345
    python gen.py "your prompt" --count 3

Output saved to ~/Pictures/mcp-generated/
"""

import argparse
import json
import random
import sys
import time
import urllib.request
from pathlib import Path

COMFYUI = "http://localhost:8188"
OUTPUT_DIR = Path.home() / "Pictures" / "mcp-generated"
WORKFLOW_PATH = Path(__file__).parent / "src/workflows/flux2_klein_heretic.json"


def load_workflow():
    with open(WORKFLOW_PATH) as f:
        return json.load(f)


def submit(workflow):
    data = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(
        f"{COMFYUI}/prompt", data=data,
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())["prompt_id"]


def wait(prompt_id, timeout=300):
    print("  ⏳ Waiting for ComfyUI...", end="", flush=True)
    start = time.time()
    while time.time() - start < timeout:
        with urllib.request.urlopen(f"{COMFYUI}/history/{prompt_id}") as resp:
            history = json.loads(resp.read())
        if prompt_id in history:
            print(" done.", flush=True)
            outputs = history[prompt_id].get("outputs", {})
            for node_out in outputs.values():
                if "images" in node_out:
                    return node_out["images"][0]
            return None
        print(".", end="", flush=True)
        time.sleep(2)
    raise TimeoutError(f"Timed out after {timeout}s")


def download(filename, subfolder=""):
    url = f"{COMFYUI}/view?filename={filename}&subfolder={subfolder}&type=output"
    with urllib.request.urlopen(url) as resp:
        return resp.read()


def generate(prompt, steps=20, width=1024, height=1024, seed=-1, name="cli"):
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)

    workflow = load_workflow()

    # Patch positive prompt (node 2)
    workflow["2"]["inputs"]["text"] = prompt
    # Patch negative prompt (node 3) — leave empty
    workflow["3"]["inputs"]["text"] = ""
    # Patch seed (node 10)
    if "10" in workflow:
        workflow["10"]["inputs"]["noise_seed"] = seed
    # Patch dimensions (node 6)
    workflow["6"]["inputs"]["width"] = width
    workflow["6"]["inputs"]["height"] = height
    # Patch steps (node 7)
    workflow["7"]["inputs"]["steps"] = steps
    # Patch output filename (node 13)
    workflow["13"]["inputs"]["filename_prefix"] = name

    print(f"  Prompt : {prompt[:80]}{'...' if len(prompt) > 80 else ''}")
    print(f"  Size   : {width}×{height}  Steps: {steps}  Seed: {seed}")

    prompt_id = submit(workflow)
    image_info = wait(prompt_id)

    if not image_info:
        print("  ❌ No output image returned.", file=sys.stderr)
        return None

    img_data = download(image_info["filename"], image_info.get("subfolder", ""))
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    out_path = OUTPUT_DIR / f"{name}_{seed}.png"
    out_path.write_bytes(img_data)
    print(f"  ✅ Saved: {out_path}  ({len(img_data) // 1024}KB)")
    return out_path


def main():
    parser = argparse.ArgumentParser(
        description="Generate images via ComfyUI FLUX.2 Klein Heretic"
    )
    parser.add_argument("prompt", help="Text prompt for the image")
    parser.add_argument("--steps", type=int, default=20, help="Inference steps (default: 20)")
    parser.add_argument("--width", type=int, default=1024, help="Width in pixels (default: 1024)")
    parser.add_argument("--height", type=int, default=1024, help="Height in pixels (default: 1024)")
    parser.add_argument("--seed", type=int, default=-1, help="Seed (-1 = random)")
    parser.add_argument("--count", type=int, default=1, help="Number of images (default: 1)")
    parser.add_argument("--name", default="cli", help="Output filename prefix (default: cli)")
    args = parser.parse_args()

    for i in range(args.count):
        if args.count > 1:
            print(f"\n[{i+1}/{args.count}]")
        seed = args.seed if args.seed != -1 else -1
        generate(
            prompt=args.prompt,
            steps=args.steps,
            width=args.width,
            height=args.height,
            seed=seed,
            name=f"{args.name}_{i+1:02d}" if args.count > 1 else args.name,
        )


if __name__ == "__main__":
    main()