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", "flux_ckpt_name") or "flux1-schnell-fp8.safetensors"
custom_suffix = jingrow.db.get_single_value("Ai Settings", "flux_custom_suffix") or "fx"
sampler_name = jingrow.db.get_single_value("Ai Settings", "flux_sampler_name") or "euler"
scheduler = jingrow.db.get_single_value("Ai Settings", "flux_scheduler") or "normal"
steps = int(jingrow.db.get_single_value("Ai Settings", "flux_steps") or 4)
cfg = int(jingrow.db.get_single_value("Ai Settings", "flux_cfg") or 1)
denoise = float(jingrow.db.get_single_value("Ai Settings", "flux_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 "blur, low quality, low resolution, watermark, cropped, out of frame, overexposed, underexposed, bad anatomy, deformed body, extra limbs, missing limbs, extra fingers, missing fingers, poor lighting, noisy background, cluttered background, blurry background, pixelated, distorted, unnatural pose, double exposure, incorrect perspective, grainy, dark, bad proportions, out of focus, incorrect clothing fit, unrealistic shadows, weird colors, cartoonish, low detail, jpeg artifacts, unnatural lighting"
add_title_to_image_name = jingrow.db.get_single_value("Ai Settings", "add_title_to_image_name") or 0

# 定义基础工作流 JSON 模板
WORKFLOW_TEMPLATE = """
{
    "3": {
        "class_type": "KSampler",
        "inputs": {
            "cfg": %d,
            "denoise": %d,
            "latent_image": [
                "5",
                0
            ],
            "model": [
                "4",
                0
            ],
            "negative": [
                "7",
                0
            ],
            "positive": [
                "6",
                0
            ],
            "sampler_name": "%s",
            "scheduler": "%s",
            "seed": 8566257,
            "steps": %d
        }
    },
    "4": {
        "class_type": "CheckpointLoaderSimple",
        "inputs": {
            "ckpt_name": "%s"
        }
    },
    "5": {
        "class_type": "EmptyLatentImage",
        "inputs": {
            "batch_size": 1,
            "height": %d,
            "width": %d
        }
    },
    "6": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "4",
                1
            ],
            "text": "masterpiece best quality girl"
        }
    },
    "7": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": [
                "4",
                1
            ],
            "text": "%s"
        }
    },
    "8": {
        "class_type": "VAEDecode",
        "inputs": {
            "samples": [
                "3",
                0
            ],
            "vae": [
                "4",
                2
            ]
        }
    },
    "save_image_websocket_node": {
        "class_type": "SaveImageWebsocket",
        "inputs": {
            "images": [
                "8",
                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 process_comfyui_flux_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:
            # 连接WebSocket
            ws.connect(f"ws://{comfyui_server_address}/ws?clientId={client_id}")
            
            for i in range(images_per_prompt):
                # 准备工作流程
                workflow = json.loads(WORKFLOW_TEMPLATE % (
                    cfg,
                    denoise,
                    sampler_name,
                    scheduler,
                    steps,
                    ckpt_name,
                    image_height,
                    image_width,
                    negative_prompt
                ))
                
                # 设置提示词
                workflow["6"]["inputs"]["text"] = prompts
                
                # 生成随机种子
                seed = random.randint(1, 4294967295)
                workflow["3"]["inputs"]["seed"] = seed
                
                # 准备文件名
                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}_"
                    
                # 只有当add_title_to_image_name为真时，才把title添加到文件名中
                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:
                    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] + "..."
            
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False

@jingrow.whitelist()
def call_jingrow_flux_api(prompts: str, config: Dict = None) -> Dict:
    """Flux 模型生成图片的 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,
            "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):
                # 准备工作流程
                workflow = json.loads(WORKFLOW_TEMPLATE % (
                    config_dict['cfg'],
                    config_dict['denoise'],
                    config_dict['sampler_name'],
                    config_dict['scheduler'],
                    config_dict['steps'],
                    config_dict['ckpt_name'],
                    config_dict['image_height'],
                    config_dict['image_width'],
                    config_dict['negative_prompt']
                ))
                
                # 设置提示词
                workflow["6"]["inputs"]["text"] = prompts
                
                # 生成随机种子
                seed = random.randint(1, 4294967295)
                workflow["3"]["inputs"]["seed"] = seed
                
                # 获取生成的图像
                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"Flux 生成图片失败: {str(e)}"
            }
            
        finally:
            if ws:
                ws.close()
                
    except Exception as e:
        return {
            "success": False,
            "message": f"调用 Flux API 异常：{str(e)}"
        }