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

import jingrow
from jingrow.ai.api.translate import get_translated_prompts
from jingrow.ai.utils.utils import update_status_progress, update_progress, send_notification, upload_image

# 定义配置变量
comfyui_server_address = jingrow.db.get_single_value("Ai Settings", "comfyui_server_address")
ckpt_name = jingrow.db.get_single_value("Ai Settings", "sd_ckpt_name") or "sd3.5_large.safetensors"
clip_l_name = jingrow.db.get_single_value("Ai Settings", "clip_l_name") or "clip_l.safetensors"
clip_g_name = jingrow.db.get_single_value("Ai Settings", "clip_g_name") or "clip_g.safetensors"
t5_name = jingrow.db.get_single_value("Ai Settings", "t5_name") or "t5xxl_fp16.safetensors"
custom_suffix = jingrow.db.get_single_value("Ai Settings", "sd_custom_suffix") or "sd3"
sampler_name = jingrow.db.get_single_value("Ai Settings", "sd_sampler_name") or "euler"
scheduler = jingrow.db.get_single_value("Ai Settings", "sd_scheduler") or "sgm_uniform"
steps = int(jingrow.db.get_single_value("Ai Settings", "sd_steps") or 30)
cfg = float(jingrow.db.get_single_value("Ai Settings", "sd_cfg") or 5.5)
denoise = float(jingrow.db.get_single_value("Ai Settings", "sd_denoise") or 1)
images_per_prompt = int(jingrow.db.get_single_value("Ai Settings", "images_per_prompt") or 4)
image_width = int(jingrow.db.get_single_value("Ai Settings", "image_width") or 1024)
image_height = int(jingrow.db.get_single_value("Ai Settings", "image_height") or 1024)
negative_prompt = jingrow.db.get_single_value("Ai Settings", "negative_prompt") or ""
add_title_to_image_name = jingrow.db.get_single_value("Ai Settings", "add_title_to_image_name") or 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": "EmptySD3LatentImage",
        "inputs": {
            "width": %d,
            "height": %d,
            "batch_size": 1
        }
    },
    "16": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "43",
                0
            ],
            "text": ""
        }
    },
    "40": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "43",
                0
            ],
            "text": ""
        }
    },
    "3": {
        "class_type": "KSampler",
        "inputs": {
            "model": [
                "4",
                0
            ],
            "positive": [
                "16",
                0
            ],
            "negative": [
                "40",
                0
            ],
            "latent_image": [
                "53",
                0
            ],
            "seed": %d,
            "steps": %d,
            "cfg": %.2f,
            "sampler_name": "%s",
            "scheduler": "%s",
            "denoise": %.2f
        }
    },
    "8": {
        "class_type": "VAEDecode",
        "inputs": {
            "samples": [
                "3",
                0
            ],
            "vae": [
                "4",
                2
            ]
        }
    },
    "9": {
        "class_type": "SaveImageWebsocket",
        "inputs": {
            "images": [
                "8",
                0
            ]
        }
    }
}
"""

def queue_prompt(prompt: Dict, comfyui_server_address: str, client_id: str) -> Dict:
    try:
        p = {"prompt": prompt, "client_id": client_id}
        data = json.dumps(p).encode('utf-8')
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        req = urllib.request.Request(
            f"http://{comfyui_server_address}/prompt",
            data=data,
            headers=headers
        )
        
        try:
            response = urllib.request.urlopen(req)
            return json.loads(response.read())
        except urllib.error.HTTPError as e:
            jingrow.log_error("ComfyUI Error", f"HTTP Error {e.code}: {e.reason}")
            raise
    except Exception as e:
        jingrow.log_error("ComfyUI Error", f"Failed to queue prompt: {str(e)}")
        raise

def get_images(ws: websocket.WebSocket, workflow: Dict, comfyui_server_address: str, client_id: str) -> Dict:
    try:
        # 确保工作流中的所有节点都有正确的格式
        for node_id, node_data in workflow.items():
            if "inputs" not in node_data:
                node_data["inputs"] = {}
            if "class_type" not in node_data:
                raise ValueError(f"Node {node_id} missing class_type")
        
        prompt_response = queue_prompt(workflow, comfyui_server_address, client_id)
        
        if not isinstance(prompt_response, dict):
            return {}
            
        prompt_id = prompt_response.get('prompt_id')
        if not prompt_id:
            return {}

    except Exception as e:
        jingrow.log_error("ComfyUI Error", f"Failed to get prompt_id: {str(e)}")
        return {}

    output_images = {}
    current_node = ""
    try:
        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 == '9':  # SaveImageWebsocket节点ID
                    images_output = output_images.get(current_node, [])
                    images_output.append(out[8:])
                    output_images[current_node] = images_output

    except Exception as e:
        jingrow.log_error("ComfyUI Error", f"Error in websocket communication: {str(e)}")
        return {}

    return output_images

def process_comfyui_sd_task(pg_name: str, pagetype: str) -> bool:
    try:
        # 获取页面对象
        pg = jingrow.get_pg(pagetype, pg_name)
        if not pg:
            return False
        
        # 从页面对象获取提示词
        prompts = None
        if hasattr(pg, "prompts") and pg.prompts:
            prompts = pg.prompts
        elif hasattr(pg, "prompts_zh") and pg.prompts_zh:
            prompts = get_translated_prompts(pg.prompts_zh)
        
        if not prompts:
            if hasattr(pg, "status"):
                pg.status = "任务未完成"
                pg.save()
                jingrow.db.commit()
            send_notification("提示词未设置，请检查配置")
            return False
        
        # 更新状态为"任务进行中"
        update_status_progress(pagetype, pg_name, "任务进行中", 20)
        
        # 创建WebSocket连接
        ws = websocket.WebSocket()
        client_id = str(uuid.uuid4())
        
        try:
            ws.connect(f"ws://{comfyui_server_address}/ws?clientId={client_id}")
            
            for i in range(images_per_prompt):
                # 生成随机种子
                seed = random.randint(1, 4294967295)
                
                # 准备工作流程
                try:
                    params = (
                        ckpt_name,
                        clip_l_name,
                        clip_g_name,
                        t5_name,
                        image_width,
                        image_height,
                        seed,
                        steps,
                        cfg,
                        sampler_name,
                        scheduler,
                        denoise
                    )
                    
                    # 格式化工作流模板
                    workflow_str = WORKFLOW_TEMPLATE % params
                    workflow = json.loads(workflow_str)
                    
                    # 设置提示词
                    workflow["16"]["inputs"]["text"] = prompts
                    workflow["40"]["inputs"]["text"] = negative_prompt
                    
                    # 移除所有空字段
                    for node in workflow.values():
                        if "widgets_values" in node:
                            del node["widgets_values"]
                    
                except Exception as e:
                    jingrow.log_error("ComfyUI Error", f"Failed to prepare workflow: {str(e)}")
                    raise

                # 准备文件名
                unique_id = str(uuid.uuid4())[:4]
                fields_to_include = ""
                
                if hasattr(pg, "serial_number") and pg.serial_number:
                    fields_to_include += f"{pg.serial_number}_"
                if int(add_title_to_image_name) and hasattr(pg, "title") and pg.title:
                    fields_to_include += f"{pg.title}_"
                    
                image_filename = f"{fields_to_include}{pg_name}_{custom_suffix}_{unique_id}_{i+1}.png"
                
                try:
                    # 获取生成的图像
                    images = get_images(ws, workflow, comfyui_server_address, client_id)
                    
                    # 上传图像
                    for node_id, image_list in images.items():
                        for image_data in image_list:
                            upload_image(pg_name, image_data, image_filename, pagetype)
                    
                    # 更新进度
                    if hasattr(pg, "progress"):
                        current_progress = 20 + ((i + 1) * 80) // images_per_prompt
                        update_progress(pg_name, min(current_progress, 100), pagetype)
                    
                except Exception as e:
                    update_status_progress(pagetype, pg_name, "任务未完成", 0)
                    return False
                    
                time.sleep(2)
        finally:
            ws.close()
        
        # 更新状态为"已出图"
        update_status_progress(pagetype, pg_name, "已出图", 100)
        return True
        
    except Exception as e:
        error_message = str(e)
        if len(error_message) > 100:
            error_message = error_message[:100] + "..."
        
        jingrow.log_error("ComfyUI Error", f"Task failed: {error_message}")
        
        try:
            pg = jingrow.get_pg(pagetype, pg_name)
            if pg and hasattr(pg, "status"):
                pg.status = "任务未完成"
                if hasattr(pg, "progress"):
                    pg.progress = 0
                pg.save()
                jingrow.db.commit()
        except Exception as inner_e:
            jingrow.log_error("ComfyUI Error", f"Failed to update error status: {str(inner_e)}")
            
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False

@jingrow.whitelist()
def call_jingrow_sd_api(prompts: str, config: Dict = None) -> Dict:
    """Stable Diffusion 模型生成图片的 API 函数，适合在 Ai Agent 节点中调用"""
    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
        }
        if config:
            config_dict.update(config)
        
        # 创建 WebSocket 连接
        ws = websocket.WebSocket()
        client_id = str(uuid.uuid4())
        
        try:
            # 连接 WebSocket
            ws.connect(f"ws://{comfyui_server_address}/ws?clientId={client_id}")
            
            # 获取需要生成的图片数量
            images_count = int(config_dict.get('images_per_prompt', 1))
            images_data = []
            
            # 循环生成指定数量的图片
            for i in range(images_count):
                # 生成随机种子
                seed = random.randint(1, 4294967295)
                
                # 准备工作流程
                try:
                    params = (
                        config_dict['ckpt_name'],
                        config_dict['clip_l_name'],
                        config_dict['clip_g_name'],
                        config_dict['t5_name'],
                        config_dict['image_width'],
                        config_dict['image_height'],
                        seed,
                        config_dict['steps'],
                        config_dict['cfg'],
                        config_dict['sampler_name'],
                        config_dict['scheduler'],
                        config_dict['denoise']
                    )
                    
                    # 格式化工作流模板
                    workflow_str = WORKFLOW_TEMPLATE % params
                    workflow = json.loads(workflow_str)
                    
                    # 设置提示词
                    workflow["16"]["inputs"]["text"] = prompts
                    workflow["40"]["inputs"]["text"] = config_dict['negative_prompt']
                    
                    # 移除所有空字段
                    for node in workflow.values():
                        if "widgets_values" in node:
                            del node["widgets_values"]
                    
                except Exception as e:
                    return {
                        "success": False,
                        "message": f"准备工作流程失败: {str(e)}"
                    }
                
                # 获取生成的图像
                images_dict = get_images(ws, workflow, comfyui_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)}"
        }