import jingrow
import json
import base64
from jingrow.utils.jinja import render_template
from jingrow.ai.api.jart import call_jingrow_jart_api
from jingrow.ai.api.jart_v1 import call_jingrow_jart_v1_api
from jingrow.ai.api.midjourney import process_midjourney_node_task
from jingrow.ai.api.comfyui_flux import call_jingrow_flux_api
from jingrow.ai.api.comfyui_sd import call_jingrow_sd_api

def build_model_config(image_tool, config):
    """构建各模型的配置参数"""
    base_config = {
        "images_per_prompt": int(config.get("images_per_prompt", 4)),
        "image_width": int(config.get("image_width", 1024)),
        "image_height": int(config.get("image_height", 1024))
    }
    
    if image_tool == "Jart":
        return base_config
        
    elif image_tool == "Jart v1":
        return base_config
        
    elif image_tool == "Flux":
        flux_config = base_config.copy()
        flux_config.update({
            "server_address": config.get("flux_server_address", ""),
            "ckpt_name": config.get("flux_ckpt_name", "flux1-schnell-fp8.safetensors"),
            "steps": int(config.get("flux_steps", 4)),
            "cfg": float(config.get("flux_cfg", 1)),
            "sampler_name": config.get("flux_sampler_name", "euler"),
            "scheduler": config.get("flux_scheduler", "normal"),
            "denoise": float(config.get("flux_denoise", 1)),
            "custom_suffix": config.get("flux_custom_suffix", "fx"),
            "negative_prompt": config.get("flux_negative_prompt", "")
        })
        return flux_config
        
    elif image_tool == "Stable Diffusion":
        sd_config = base_config.copy()
        sd_config.update({
            "server_address": config.get("sd_server_address", ""),
            "ckpt_name": config.get("sd_ckpt_name", "sd3.5_large.safetensors"),
            "clip_l_name": config.get("clip_l_name", "clip_l.safetensors"),
            "clip_g_name": config.get("clip_g_name", "clip_g.safetensors"),
            "t5_name": config.get("t5_name", "t5xxl_fp16.safetensors"),
            "steps": int(config.get("sd_steps", 30)),
            "cfg": float(config.get("sd_cfg", 5.5)),
            "sampler_name": config.get("sd_sampler_name", "euler"),
            "scheduler": config.get("sd_scheduler", "sgm_uniform"),
            "denoise": float(config.get("sd_denoise", 1)),
            "custom_suffix": config.get("sd_custom_suffix", "sd"),
            "negative_prompt": config.get("sd_negative_prompt", "")
        })
        return sd_config
        
    elif image_tool == "Midjourney":
        # 构建Midjourney选项参数
        mj_options = {}
        
        # 基础参数映射
        if config.get("version"):
            mj_options["v"] = config.get("version")
        if config.get("aspect_ratio"):
            mj_options["ar"] = config.get("aspect_ratio")
        if config.get("quality"):
            mj_options["q"] = config.get("quality")
        if config.get("raw_mode"):
            mj_options["raw"] = True
        if config.get("chaos"):
            mj_options["c"] = int(config.get("chaos"))
        if config.get("weird"):
            mj_options["weird"] = int(config.get("weird"))
        if config.get("stylize"):
            mj_options["stylize"] = int(config.get("stylize"))
        if config.get("mj_seed"):
            mj_options["seed"] = int(config.get("mj_seed"))
        if config.get("tile"):
            mj_options["tile"] = True
        # 权重参数映射 - 只在有对应参考图片时才添加权重
        # 图片权重：当有参考图片时（对直接放在prompt前面的图片有效）
        if config.get("image_weight") and config.get("reference_image"):
            mj_options["iw"] = config.get("image_weight")
        # 全能参考权重：当有全能参考时（对--ow参数有效）
        if config.get("omni_weight") and config.get("omni_reference"):
            mj_options["ow"] = config.get("omni_weight")
        # 风格权重：当有风格参考时（对--sref参数有效）
        if config.get("style_weight") and config.get("style_reference"):
            mj_options["sw"] = config.get("style_weight")
        # 角色权重：当有角色参考时（对--cref参数有效）
        if config.get("character_weight") and config.get("character_reference"):
            mj_options["cw"] = config.get("character_weight")
            
        # 构建配置对象
        mj_config = {
            "options": mj_options
        }
        
        # 处理参考图片 - 根据Midjourney API的sources参数结构
        sources = {}
        
        def parse_multiline_urls(text):
            """解析多行文本框中的URL，每行一个URL"""
            if not text:
                return []
            # 按行分割，去除空行和空白字符
            urls = [line.strip() for line in text.split('\n') if line.strip()]
            return urls
        
        # 处理参考图片（作为images）
        image_urls = []
        if config.get("reference_image"):
            image_urls.extend(parse_multiline_urls(config.get("reference_image")))
        if image_urls:
            sources["images"] = image_urls
            
        # 处理全能参考（作为独立的omni参数）
        if config.get("omni_reference"):
            omni_url = config.get("omni_reference").strip()
            if omni_url:
                sources["omni"] = [omni_url]
            
        # 处理风格参考（作为styles）
        style_urls = []
        if config.get("style_reference"):
            style_urls.extend(parse_multiline_urls(config.get("style_reference")))
        if style_urls:
            sources["styles"] = style_urls
            
        # 处理角色参考（作为caracters）
        if config.get("character_reference"):
            character_url = config.get("character_reference").strip()
            if character_url:
                sources["caracters"] = [character_url]
            
        if sources:
            mj_config["sources"] = sources
            
        # 处理负面提示词
        if config.get("no_text"):
            mj_config["negative_prompt"] = config.get("no_text")
            
        return mj_config
    
    return base_config

@jingrow.whitelist()
def execute(context=None, inputs=None, config=None, **kwargs):
    if context is None:
        context = kwargs.get("context", {})
    if inputs is None:
        inputs = kwargs.get("inputs", {})
    if config is None:
        config = kwargs.get("config", {})

    # 获取同一分支的所有上游节点数据并平铺到 inputs
    current_node_id = context.get("current_node_id")
    flow_data = context.get("flow_data", {})
    branch_upstream_nodes = get_branch_upstream_nodes(current_node_id, flow_data)
    
    if context and "node_results" in context and branch_upstream_nodes:
        node_results = context["node_results"]
        for upstream_node_id in branch_upstream_nodes:
            if upstream_node_id in node_results:
                upstream_result = node_results[upstream_node_id]
                if isinstance(upstream_result, dict):
                    # 如果有ai_raw_response，尝试解析JSON内容
                    if "ai_raw_response" in upstream_result and isinstance(upstream_result["ai_raw_response"], str):
                        try:
                            raw_response = upstream_result["ai_raw_response"].strip()
                            start_brace = raw_response.find('{')
                            end_brace = raw_response.rfind('}')
                            
                            if start_brace != -1 and end_brace != -1 and end_brace > start_brace:
                                json_content = raw_response[start_brace:end_brace + 1]
                                parsed_json = json.loads(json_content)
                                if isinstance(parsed_json, dict):
                                    for k, v in parsed_json.items():
                                        if k not in inputs or not inputs.get(k):
                                            inputs[k] = v
                        except:
                            continue
                    
                    # 平铺其他字段到inputs
                    for k, v in upstream_result.items():
                        if k not in inputs or not inputs.get(k):
                            inputs[k] = v

    # 1. 获取输入对象
    if not isinstance(inputs, dict):
        try:
            inputs = json.loads(inputs)
        except Exception:
            inputs = {}

    # 2. 获取配置
    image_tool = config.get("image_tool", "Jart")
    prompts = config.get("prompts", "")  # 前端传递的提示词

    # 3. 构建模板上下文
    template_context = inputs.copy()
    
    # 根据flow_id获取完整的agent对象
    agent_data = {}
    if context and "flow_id" in context and context["flow_id"]:
        try:
            flow_id = context["flow_id"]
            agent = jingrow.get_pg("Ai Agent", flow_id)
            if agent:
                agent_data = agent.as_dict()
        except Exception as e:
            jingrow.log_error("获取agent数据失败", f"flow_id: {flow_id}, error: {str(e)}")
            agent_data = {}
    
    # 将agent的所有字段直接平铺到模板上下文
    for key, value in agent_data.items():
        if key not in template_context or not template_context.get(key):
            template_context[key] = value

    # 4. 获取提示词
    if not prompts:
        return {
            "success": False,
            "error": "提示词不能为空，请在节点配置中输入提示词"
        }
    
    try:
        final_prompts = render_template(prompts, template_context)
    except Exception as e:
        return {
            "success": False,
            "error": f"提示词模板渲染失败: {str(e)}"
        }
    


    if not final_prompts:
        return {
            "success": False,
            "error": "提示词不能为空"
        }

    # 5. 调用对应的绘画API
    try:
        # 构建模型配置参数
        model_config = build_model_config(image_tool, config)
        
        # 根据不同的绘画工具调用对应的API函数
        if image_tool == "Jart":
            result = call_jingrow_jart_api([final_prompts], config=model_config)
            
        elif image_tool == "Jart v1":
            result = call_jingrow_jart_v1_api([final_prompts], config=model_config)
            
        elif image_tool == "Flux":
            result = call_jingrow_flux_api(final_prompts, config=model_config)
            
        elif image_tool == "Stable Diffusion":
            result = call_jingrow_sd_api(final_prompts, config=model_config)
            
        elif image_tool == "Midjourney":
            result = process_midjourney_node_task(final_prompts, config=model_config)
            
        else:
            return {
                "success": False,
                "error": f"不支持的绘画工具: {image_tool}"
            }

        # 6. 处理结果
        if result.get("success"):
            # 获取生成的图像数据
            image_data = result.get("data", [])
            
            if image_data:
                # 将bytes数据转换为base64编码，确保可以正常序列化
                encoded_image_data = []
                for img_item in image_data:
                    if isinstance(img_item, dict) and "image_data" in img_item:
                        original_data = img_item["image_data"]
                        if isinstance(original_data, bytes):
                            # 将bytes转换为base64字符串
                            base64_data = base64.b64encode(original_data).decode('utf-8')
                            encoded_img_item = {
                                "image_data": base64_data,
                                "data_type": "base64",  # 标记数据类型
                                "original_size": len(original_data)  # 记录原始大小
                            }
                            # 保留其他字段
                            for key, value in img_item.items():
                                if key != "image_data":
                                    encoded_img_item[key] = value
                            encoded_image_data.append(encoded_img_item)
                        else:
                            # 如果不是bytes，直接保留
                            encoded_image_data.append(img_item)
                    else:
                        # 如果不是预期的格式，直接保留
                        encoded_image_data.append(img_item)
                
                # 返回编码后的图片数据，供后续节点使用
                return {
                    "success": True,
                    "images_count": len(encoded_image_data),
                    "image_data": encoded_image_data
                }
            else:
                return {
                    "success": False,
                    "error": "未生成任何图像"
                }
        else:
            error_message = result.get("message", "未知错误")
            return {
                "success": False,
                "error": f"{image_tool} 生成失败: {error_message}"
            }

    except Exception as e:
        return {
            "success": False,
            "error": f"调用 {image_tool} API 失败: {str(e)}"
        }

def get_branch_upstream_nodes(current_node_id, flow_data):
    """获取同一分支的所有上游节点ID列表"""
    if not current_node_id or not flow_data:
        return []
    
    edges = flow_data.get("edges", [])
    upstream_nodes = set()
    stack = [current_node_id]
    
    while stack:
        node_id = stack.pop()
        for edge in edges:
            if edge.get("target") == node_id:
                source_id = edge.get("source")
                if source_id and source_id not in upstream_nodes:
                    upstream_nodes.add(source_id)
                    stack.append(source_id)
    
    return list(upstream_nodes) 