#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 == '94':
                images_output = output_images.get(current_node, [])
                images_output.append(out[8:])
                output_images[current_node] = images_output

    return output_images

prompt_text = """
{
  "76": {
    "inputs": {
      "value": "A hedgehog wearing a tiny party hat surrounded by confetti, early digital camera style, slight noise, flash photography, candid moment, 2000s digicam aesthetic, festive birthday celebration atmosphere\n"
    },
    "class_type": "PrimitiveStringMultiline",
    "_meta": {
      "title": "Prompt"
    }
  },
  "94": {
    "inputs": {
      "images": [
        "77:82",
        0
      ]
    },
    "class_type": "SaveImageWebsocket",
    "_meta": {
      "title": "Save Image (Websocket)"
    }
  },
  "77:80": {
    "inputs": {
      "sampler_name": "euler"
    },
    "class_type": "KSamplerSelect",
    "_meta": {
      "title": "KSamplerSelect"
    }
  },
  "77:81": {
    "inputs": {
      "noise": [
        "77:86",
        0
      ],
      "guider": [
        "77:90",
        0
      ],
      "sampler": [
        "77:80",
        0
      ],
      "sigmas": [
        "77:93",
        0
      ],
      "latent_image": [
        "77:83",
        0
      ]
    },
    "class_type": "SamplerCustomAdvanced",
    "_meta": {
      "title": "SamplerCustomAdvanced"
    }
  },
  "77:82": {
    "inputs": {
      "samples": [
        "77:81",
        0
      ],
      "vae": [
        "77:89",
        0
      ]
    },
    "class_type": "VAEDecode",
    "_meta": {
      "title": "VAE Decode"
    }
  },
  "77:83": {
    "inputs": {
      "width": [
        "77:84",
        0
      ],
      "height": [
        "77:85",
        0
      ],
      "batch_size": 1
    },
    "class_type": "EmptyFlux2LatentImage",
    "_meta": {
      "title": "Empty Flux 2 Latent"
    }
  },
  "77:84": {
    "inputs": {
      "value": 1024
    },
    "class_type": "PrimitiveInt",
    "_meta": {
      "title": "Width"
    }
  },
  "77:85": {
    "inputs": {
      "value": 1024
    },
    "class_type": "PrimitiveInt",
    "_meta": {
      "title": "Height"
    }
  },
  "77:86": {
    "inputs": {
      "noise_seed": 432262096973490
    },
    "class_type": "RandomNoise",
    "_meta": {
      "title": "RandomNoise"
    }
  },
  "77:87": {
    "inputs": {
      "unet_name": "flux-2-klein-4b.safetensors",
      "weight_dtype": "default"
    },
    "class_type": "UNETLoader",
    "_meta": {
      "title": "Load Diffusion Model"
    }
  },
  "77:88": {
    "inputs": {
      "clip_name": "qwen_3_4b.safetensors",
      "type": "flux2",
      "device": "default"
    },
    "class_type": "CLIPLoader",
    "_meta": {
      "title": "Load CLIP"
    }
  },
  "77:89": {
    "inputs": {
      "vae_name": "flux2-vae.safetensors"
    },
    "class_type": "VAELoader",
    "_meta": {
      "title": "Load VAE"
    }
  },
  "77:90": {
    "inputs": {
      "cfg": 1,
      "model": [
        "77:87",
        0
      ],
      "positive": [
        "77:92",
        0
      ],
      "negative": [
        "77:91",
        0
      ]
    },
    "class_type": "CFGGuider",
    "_meta": {
      "title": "CFG Guider"
    }
  },
  "77:91": {
    "inputs": {
      "conditioning": [
        "77:92",
        0
      ]
    },
    "class_type": "ConditioningZeroOut",
    "_meta": {
      "title": "ConditioningZeroOut"
    }
  },
  "77:92": {
    "inputs": {
      "text": [
        "76",
        0
      ],
      "clip": [
        "77:88",
        0
      ]
    },
    "class_type": "CLIPTextEncode",
    "_meta": {
      "title": "CLIP Text Encode (Positive Prompt)"
    }
  },
  "77:93": {
    "inputs": {
      "steps": 4,
      "width": [
        "77:84",
        0
      ],
      "height": [
        "77:85",
        0
      ]
    },
    "class_type": "Flux2Scheduler",
    "_meta": {
      "title": "Flux2Scheduler"
    }
  }
}
"""

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="A Linux Laptop on a table showing on the screen: PHORONIX is great!")
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=1024, help="Image width (default: 1024, max: 2048)")
parser.add_argument("--height", type=int, default=1024, help="Image height (default: 1024, max: 2048)")
parser.add_argument("--9b", "--9b-model", dest="use_9b", action="store_true", help="Use 9B model instead of 4B")
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 = json.loads(' '.join(prompt_text.split()))
prompt["76"]["inputs"]["value"] = args.prompt

if args.use_9b:
    prompt["77:87"]["inputs"]["unet_name"] = "flux-2-klein-9b-fp8.safetensors"
    prompt["77:88"]["inputs"]["clip_name"] = "qwen_3_8b_fp8mixed.safetensors"

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["77:84"]["inputs"]["value"] = width
prompt["77:85"]["inputs"]["value"] = 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["77:86"]["inputs"]["noise_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")