import jingrow
import json
import uuid
import base64
import requests

@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)

    # 只平铺同一分支上游节点的结果到 inputs
    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. 获取配置
    save_location = config.get("save_location", "file")
    target_pagetype = config.get("target_pagetype", "")
    target_name = config.get("target_name", "")
    filename_prefix = config.get("filename_prefix", "AI_Generated")
    save_field = config.get("save_field", "")

    # 3. 获取上游节点传递的图片数据
    image_data = inputs.get("image_data", [])
    
    if not image_data:
        error_msg = "未找到图片数据，请确保上游节点已生成图片"
        return {
            "success": False,
            "error": error_msg
        }

    # 4. 获取目标记录信息
    target_record_info = {}
    if save_location == "auto":
        # 自动模式：尝试从上游节点结果中获取记录信息
        
        # 方法1：从 inputs 中直接获取
        if "pagetype" in inputs and "name" in inputs:
            try:
                target_pg = jingrow.get_pg(inputs["pagetype"], inputs["name"])
                if target_pg:
                    target_record_info = {
                        "pagetype": inputs["pagetype"],
                        "name": inputs["name"],
                        "pg": target_pg
                    }
            except Exception:
                pass
        
        # 方法2：从 context.node_results 中查找包含记录信息的节点
        if not target_record_info and context and "node_results" in context:
            node_results = context["node_results"]
            
            # 查找包含记录信息的节点结果
            for node_id, node_result in node_results.items():
                if isinstance(node_result, dict):
                    # 检查是否包含记录信息
                    if "pagetype" in node_result and "name" in node_result:
                        try:
                            target_pg = jingrow.get_pg(node_result["pagetype"], node_result["name"])
                            if target_pg:
                                target_record_info = {
                                    "pagetype": node_result["pagetype"],
                                    "name": node_result["name"],
                                    "pg": target_pg
                                }
                                break
                        except Exception:
                            continue
            
            # 如果还没找到，尝试从其他字段中查找
            if not target_record_info:
                for node_id, node_result in node_results.items():
                    if isinstance(node_result, dict):
                        # 查找可能包含记录信息的字段
                        for field_name, field_value in node_result.items():
                            if field_name in ["record", "pg", "document"] and hasattr(field_value, "pagetype") and hasattr(field_value, "name"):
                                try:
                                    target_pg = jingrow.get_pg(field_value.pagetype, field_value.name)
                                    if target_pg:
                                        target_record_info = {
                                            "pagetype": field_value.pagetype,
                                            "name": field_value.name,
                                            "pg": target_pg
                                        }
                                        break
                                except Exception:
                                    continue
                        if target_record_info:
                            break
        
        # 如果自动模式仍然没有找到记录信息，回退到文件库模式
        if not target_record_info:
            save_location = "file"

    elif save_location == "record":
        # 优先使用配置中的目标记录信息
        if target_pagetype and target_name:
            try:
                target_pg = jingrow.get_pg(target_pagetype, target_name)
                if target_pg:
                    target_record_info = {
                        "pagetype": target_pagetype,
                        "name": target_name,
                        "pg": target_pg
                    }
                else:
                    error_msg = f"无法找到指定的目标记录: {target_pagetype} - {target_name}"
                    return {
                        "success": False,
                        "error": error_msg
                    }
            except Exception as e:
                error_msg = f"获取目标记录失败: {str(e)}"
                return {
                    "success": False,
                    "error": error_msg
                }
        # 如果没有配置，尝试从上游节点结果中获取
        elif "pagetype" in inputs and "name" in inputs:
            try:
                target_pg = jingrow.get_pg(inputs["pagetype"], inputs["name"])
                if target_pg:
                    target_record_info = {
                        "pagetype": inputs["pagetype"],
                        "name": inputs["name"],
                        "pg": target_pg
                    }
                else:
                    error_msg = f"无法找到指定的源记录: {inputs['pagetype']} - {inputs['name']}"
                    return {
                        "success": False,
                        "error": error_msg
                    }
            except Exception as e:
                error_msg = f"获取源记录失败: {str(e)}"
                return {
                    "success": False,
                    "error": error_msg
                }

    # 5. 上传图片
    failed_count = 0
    
    try:
        for img_data in image_data:
            try:
                # 检查图片数据格式
                if not isinstance(img_data, dict) or "image_data" not in img_data:
                    failed_count += 1
                    continue
                
                # 检查并处理图片数据
                raw_image_data = img_data["image_data"]
                
                # 处理base64编码的图片数据
                if isinstance(raw_image_data, str):
                    try:
                        if img_data.get("data_type") == "base64":
                            raw_image_data = base64.b64decode(raw_image_data)
                        elif raw_image_data.startswith('data:image'):
                            base64_data = raw_image_data.split(',')[1]
                            raw_image_data = base64.b64decode(base64_data)
                        elif raw_image_data.startswith('http'):
                            response = requests.get(raw_image_data, verify=False)
                            raw_image_data = response.content
                        else:
                            raw_image_data = base64.b64decode(raw_image_data)
                    except Exception:
                        failed_count += 1
                        continue
                elif isinstance(raw_image_data, bytes):
                    pass
                else:
                    failed_count += 1
                    continue
                
                if not isinstance(raw_image_data, bytes):
                    failed_count += 1
                    continue
                
                file_extension = "png"
                if "file_extension" in img_data:
                    file_extension = img_data["file_extension"]
                elif "format" in img_data:
                    file_extension = img_data["format"].lower()
                
                record_name = target_record_info.get("name", "") if target_record_info else ""
                record_prefix = f"{record_name}_" if record_name else ""
                filename = f"{filename_prefix}{record_prefix}{uuid.uuid4().hex[:8]}.{file_extension}"
                
                try:
                    attached_to_field = None
                    if save_field and target_record_info and target_record_info.get("pg"):
                        meta = jingrow.get_meta(target_record_info["pagetype"])
                        label2field = {}
                        if hasattr(meta, 'fields'):
                            for f in meta.fields:
                                if hasattr(f, 'label') and hasattr(f, 'fieldname'):
                                    label2field[str(f.label).strip()] = str(f.fieldname).strip()
                        
                        target_field = save_field
                        if save_field in label2field:
                            target_field = label2field[save_field]
                        else:
                            for label, fieldname in label2field.items():
                                try:
                                    user_lang = jingrow.get_pg("User", jingrow.session.user).language or jingrow.get_single("System Settings").language or "en"
                                    translated_label = jingrow._(label, lang=user_lang)
                                    if save_field == translated_label:
                                        target_field = fieldname
                                        break
                                except:
                                    continue
                        
                        # 只有当目标字段存在且字段值为空时，才关联到字段
                        if target_field:
                            # 检查目标字段的当前值
                            current_field_value = target_record_info["pg"].get(target_field)
                            if not current_field_value:  # 字段值为空时才关联
                                attached_to_field = target_field
                    
                    # 创建文件记录
                    if save_location == "file":
                        # 文件库模式：不关联到任何记录
                        file_data = {
                            "pagetype": "File",
                            "file_name": filename,
                            "is_private": 0,
                            "content": raw_image_data
                        }
                    else:
                        # 记录模式：关联到指定记录
                        if not target_record_info:
                            failed_count += 1
                            continue
                            
                        file_data = {
                            "pagetype": "File",
                            "attached_to_pagetype": target_record_info["pagetype"],
                            "attached_to_name": target_record_info["name"],
                            "file_name": filename,
                            "is_private": 0,
                            "content": raw_image_data
                        }
                    
                    # 如果需要关联到字段，添加 attached_to_field
                    if attached_to_field:
                        file_data["attached_to_field"] = attached_to_field
                    
                    file_pg = jingrow.get_pg(file_data)
                    file_pg.save()
                    jingrow.db.commit()
                    file_url = file_pg.file_url
                    
                    # 如果关联到字段，更新关联字段
                    if attached_to_field:
                        jingrow.db.set_value(target_record_info["pagetype"], target_record_info["name"], attached_to_field, file_url)
                    
                    if not file_url:
                        failed_count += 1
                except Exception as e:
                    # 记录具体错误信息以便调试
                    jingrow.log_error("图片上传失败", f"上传图片时发生错误: {str(e)}")
                    failed_count += 1
                    
            except Exception as e:
                # 记录具体错误信息以便调试
                jingrow.log_error("图片处理失败", f"处理图片数据时发生错误: {str(e)}")
                failed_count += 1
                continue

    except Exception as e:
        error_msg = f"图片上传过程中发生异常: {str(e)}"
        return {
            "success": False,
            "error": error_msg
        }

    # 6. 返回结果
    if failed_count > 0:
        error_msg = f"成功上传 {len(image_data) - failed_count} 张图片，失败 {failed_count} 张"
        if save_location == "file":
            error_msg += "。请检查图片数据格式是否正确，或查看系统日志获取详细错误信息。"
        else:
            error_msg += "。请检查目标记录是否存在，或查看系统日志获取详细错误信息。"
        return {
            "success": False,
            "error": error_msg
        }
    
    success_msg = f"成功上传 {len(image_data)} 张图片"
    if save_location == "file":
        success_msg += "到文件库"
    else:
        success_msg += f"到记录 {target_record_info.get('pagetype', '')} - {target_record_info.get('name', '')}"
    
    return {
        "success": True,
        "message": success_msg
    }

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)