#This is an example that uses the websockets api and the SaveImageWebsocket node to get images directly without
#them being saved to disk

import random
import websocket
import uuid
import json
import urllib.request
import urllib.parse
import argparse
import time
import os

server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())
token = ""

def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    protocol = "https" if use_https else "http"
    req =  urllib.request.Request("{}://{}/prompt?token={}".format(protocol, server_address, token), data=data)
    return json.loads(urllib.request.urlopen(req).read())

def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    protocol = "https" if use_https else "http"
    with urllib.request.urlopen("{}://{}/view?{}".format(protocol, server_address, url_values)) as response:
        return response.read()

def get_history(prompt_id):
    protocol = "https" if use_https else "http"
    with urllib.request.urlopen("{}://{}/history/{}".format(protocol, server_address, prompt_id)) as response:
        return json.loads(response.read())

def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    print(f"Prompt ID: {prompt_id}")
    output_images = {}
    current_node = ""
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            print(f"Message type: {message.get('type')}")
            if message['type'] == 'executing':
                data = message['data']
                if data['prompt_id'] == prompt_id:
                    if data['node'] is None:
                        break #Execution is done
                    else:
                        current_node = data['node']
                        print(f"Executing node: {current_node}")
        else:
            print(f"Received binary data, current_node: {current_node}")
            if current_node == '271':
                images_output = output_images.get(current_node, [])
                images_output.append(out[8:])
                output_images[current_node] = images_output

    return output_images

parser = argparse.ArgumentParser()
parser.add_argument("--url", type=str, default="http://127.0.0.1:8188", help="ComfyUI server URL (default: http://127.0.0.1:8188)")
parser.add_argument("--prompt", type=str, default=None)
parser.add_argument("--seed", type=int, default=None)
parser.add_argument("--runs", type=int, default=1, help="Number of times to run the generation")
parser.add_argument("--width", type=int, default=None, help="Image width (max: 2048)")
parser.add_argument("--height", type=int, default=None, help="Image height (max: 2048)")
parser.add_argument("--token", type=str, default=None, help="API token")
parser.add_argument("--token-file", type=str, default=None, help="File containing API token (default: token.txt)")
parser.add_argument("--output", type=str, default="output", help="Output directory (default: output)")
args = parser.parse_args()

if args.token:
    token = args.token
else:
    if args.token_file:
        token_path = args.token_file
    elif os.path.exists("token.txt"):
        token_path = "token.txt"
    else:
        print("Error: No token provided. Create token.txt or use --token or --token-file option.")
        exit(1)
    
    if not os.path.exists(token_path):
        print(f"Error: Token file '{token_path}' not found.")
        exit(1)
    
    with open(token_path) as f:
        token = f.read().strip()

os.makedirs(args.output, exist_ok=True)
server_address = args.url
use_https = server_address.startswith("https://")
if use_https:
    server_address = server_address.replace("https://", "")
    ws_protocol = "wss"
else:
    server_address = server_address.replace("http://", "") if server_address.startswith("http://") else server_address
    ws_protocol = "ws"

prompt_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qwen_image_2512_360_lora_ws.json")
with open(prompt_path) as f:
    prompt = json.load(f)

if args.prompt is None:
    args.prompt = prompt["269"]["inputs"]["value"]
if args.width is None:
    args.width = prompt["238:232"]["inputs"]["width"]
if args.height is None:
    args.height = prompt["238:232"]["inputs"]["height"]
prompt["269"]["inputs"]["value"] = args.prompt

from PIL import Image
import io
import os
from datetime import datetime

os.makedirs("output", exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

total_start_time = time.time()
total_images = 0

width = min(max(args.width, 256), 2048)
height = min(max(args.height, 256), 2048)
prompt["238:232"]["inputs"]["width"] = width
prompt["238:232"]["inputs"]["height"] = height

print(f"Generating {args.runs} image(s) at {width}x{height}")
print(f"Prompt: {args.prompt[:100]}{'...' if len(args.prompt) > 100 else ''}")

for run_idx in range(args.runs):
    run_start_time = time.time()
    if args.seed is not None:
        seed = args.seed + run_idx
    else:
        seed = random.getrandbits(64)
    prompt["238:230"]["inputs"]["seed"] = seed

    print(f"Run {run_idx + 1}/{args.runs} - Seed: {seed}")

    ws = websocket.WebSocket()
    ws.connect("{}://{}/ws?clientId={}&token={}".format(ws_protocol, server_address, client_id, token))
    images = get_images(ws, prompt)
    ws.close()

    for i, node_id in enumerate(images):
        for j, image_data in enumerate(images[node_id]):
            image = Image.open(io.BytesIO(image_data))
            out_path = f"{args.output}/image_{timestamp}_{seed}_{run_idx}_{i}_{j}.png"
            image.save(out_path)
            print(f"Saved: {out_path}")
            total_images += 1

    run_elapsed = time.time() - run_start_time
    print(f"Run {run_idx + 1} completed - Images: {len(images)}, Time: {run_elapsed:.2f}s")

total_elapsed = time.time() - total_start_time
print(f"All runs completed - Total runs: {args.runs}, Total images: {total_images}, Total time: {total_elapsed:.2f}s")