import json
import base64
import requests
import random
import time
import websocket
import uuid
import urllib.request
from typing import Dict, List, Any, Optional

# 定义配置变量
comfyui_server_address = "127.0.0.1:8188"
ckpt_name = "sd3.5_large.safetensors"
clip_l_name = "clip_l.safetensors"
clip_g_name = "clip_g.safetensors"
t5_name = "t5xxl_fp16.safetensors"
custom_suffix = "sd3"
sampler_name = "euler"
scheduler = "sgm_uniform"
steps = 30
cfg = 5.5
denoise = 1.0
images_per_prompt = 4
image_width = 1024
image_height = 1024
negative_prompt = ""
add_title_to_image_name = 0

# 定义基础工作流 JSON 模板
WORKFLOW_TEMPLATE = """
{
    "4": {
        "class_type": "CheckpointLoaderSimple",
        "inputs": {
            "ckpt_name": "%s"
        }
    },
    "43": {
        "class_type": "TripleCLIPLoader",
        "inputs": {
            "clip_name1": "%s",
            "clip_name2": "%s",
            "clip_name3": "%s"
        }
    },
    "53": {
        "class_type": "KSampler",
        "inputs": {
            "cfg": %f,
            "denoise": %f,
            "latent_image": [
                "54",
                0
            ],
            "model": [
                "4",
                0
            ],
            "negative": [
                "56",
                0
            ],
            "positive": [
                "55",
                0
            ],
            "sampler_name": "%s",
            "scheduler": "%s",
            "seed": 8566257,
            "steps": %d
        }
    },
    "54": {
        "class_type": "EmptyLatentImage",
        "inputs": {
            "batch_size": 1,
            "height": %d,
            "width": %d
        }
    },
    "55": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "43",
                0
            ],
            "text": "masterpiece best quality girl"
        }
    },
    "56": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "43",
                0
            ],
            "text": "%s"
        }
    },
    "57": {
        "class_type": "VAEDecode",
        "inputs": {
            "samples": [
                "53",
                0
            ],
            "vae": [
                "4",
                2
            ]
        }
    },
    "save_image_websocket_node": {
        "class_type": "SaveImageWebsocket",
        "inputs": {
            "images": [
                "57",
                0
            ]
        }
    }
}
"""

def queue_prompt(prompt: Dict, comfyui_server_address: str, client_id: str) -> Dict:
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request(f"http://{comfyui_server_address}/prompt", data=data)
    response = json.loads(urllib.request.urlopen(req).read())
    return response

def get_images(ws: websocket.WebSocket, workflow: Dict, comfyui_server_address: str, client_id: str) -> Dict:
    try:
        prompt_response = queue_prompt(workflow, comfyui_server_address, client_id)
        prompt_id = prompt_response['prompt_id']
    except KeyError:
        return {}

    output_images = {}
    current_node = ""
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data.get('prompt_id') == prompt_id:
                    if data['node'] is None:
                        break
                    else:
                        current_node = data['node']
        else:
            if current_node == 'save_image_websocket_node':
                images_output = output_images.get(current_node, [])
                images_output.append(out[8:])
                output_images[current_node] = images_output

    return output_images

def call_jingrow_sd_api(prompts: str, config: Dict = None) -> Dict:
    """Stable Diffusion 模型生成图片的 API 函数"""
    try:
        # 合并配置
        config_dict = {
            "images_per_prompt": images_per_prompt,
            "image_width": image_width,
            "image_height": image_height,
            "cfg": cfg,
            "denoise": denoise,
            "sampler_name": sampler_name,
            "scheduler": scheduler,
            "steps": steps,
            "ckpt_name": ckpt_name,
            "clip_l_name": clip_l_name,
            "clip_g_name": clip_g_name,
            "t5_name": t5_name,
            "negative_prompt": negative_prompt,
            "server_address": comfyui_server_address
        }
        if config:
            config_dict.update(config)
        
        # 使用配置中的服务器地址
        server_address = config_dict.get('server_address', comfyui_server_address)
        
        # 创建 WebSocket 连接
        ws = websocket.WebSocket()
        client_id = str(uuid.uuid4())
        
        try:
            # 连接 WebSocket
            ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
            
            # 获取需要生成的图片数量
            images_count = int(config_dict.get('images_per_prompt', 1))
            images_data = []
            
            # 循环生成指定数量的图片
            for i in range(images_count):
                # 准备工作流程
                workflow = json.loads(WORKFLOW_TEMPLATE % (
                    config_dict['ckpt_name'],
                    config_dict['clip_l_name'],
                    config_dict['clip_g_name'],
                    config_dict['t5_name'],
                    config_dict['cfg'],
                    config_dict['denoise'],
                    config_dict['sampler_name'],
                    config_dict['scheduler'],
                    config_dict['steps'],
                    config_dict['image_height'],
                    config_dict['image_width'],
                    config_dict['negative_prompt']
                ))
                
                # 设置提示词
                workflow["55"]["inputs"]["text"] = prompts
                
                # 生成随机种子
                seed = random.randint(1, 4294967295)
                workflow["53"]["inputs"]["seed"] = seed
                
                # 获取生成的图像
                images_dict = get_images(ws, workflow, server_address, client_id)
                
                # 将图像数据添加到结果中
                for node_id, image_list in images_dict.items():
                    for image_data in image_list:
                        images_data.append({
                            "image_data": image_data,
                            "seed": seed,
                            "index": i + 1
                        })
            
            return {
                "success": True,
                "data": images_data
            }
            
        except Exception as e:
            return {
                "success": False,
                "message": f"Stable Diffusion 生成图片失败: {str(e)}"
            }
            
        finally:
            if ws:
                ws.close()
                
    except Exception as e:
        return {
            "success": False,
            "message": f"调用 Stable Diffusion API 异常：{str(e)}"
        }