import jingrow
import json
from jingrow.utils.jinja import render_template
from jingrow.ai.utils.utils import call_ai_model  # 新增：直接调用工具函数

@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", {})
    # jingrow.log_error("ai_content_generation inputs/config/context",  f"inputs={inputs} config={config} context={context}")

    # 获取同一分支的所有上游节点数据并平铺到 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. 获取输入对象（移除pg兼容，仅支持inputs为dict）
    if not isinstance(inputs, dict):
        try:
            inputs = json.loads(inputs)
        except Exception:
            inputs = {}

    # 2. 检查模板配置
    template = config.get("ai_user_content")
    if not template:
        error_result = {
            "success": False,
            "error": "AI内容生成节点缺少模板配置(ai_user_content)"
        }
        return error_result

    # 3. 构建模板上下文，包含智能体字段
    template_context = inputs.copy()  # 复制inputs作为基础上下文
    
    # 根据flow_id获取完整的agent对象
    agent_data = {}
    
    if context and "flow_id" in context and context["flow_id"]:
        try:
            flow_id = context["flow_id"]
            # 获取完整的Ai Agent记录
            agent = jingrow.get_pg("Ai Agent", flow_id)
            if agent:
                # 将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. 渲染模板
    try:
        # 现在模板里可以直接用{{ exclude_prompts }}、{{ name }}等任何agent字段
        prompt = render_template(template, template_context)
    except Exception as e:
        error_result = {
            "success": False,
            "error": f"模板渲染失败: {str(e)}"
        }
        return error_result

    # 5. 调用AI（改为直接用工具函数）
    system_message = config.get("ai_system_message")
    text_model = config.get("text_model") or "Jingrow"
    max_tokens = config.get("max_tokens", 1000)
    temperature = config.get("temperature", 0.9)
    top_p = config.get("top_p", 0.9)

    try:
        # 直接调用工具函数
        response = call_ai_model(
            prompt,
            text_model=text_model,
            ai_temperature=temperature,
            ai_top_p=top_p,
            ai_system_message=system_message,
            max_tokens=max_tokens,
            image_urls=None
        )
        
        # jingrow.log_error("AI内容生成节点 - AI返回内容", f"完整响应: {response}")
        
    except Exception as e:
        error_result = {
            "success": False,
            "error": f"AI模型调用失败: {str(e)}"
        }
        return error_result

    # 6. 返回AI输出
    try:
        result = {
            "success": True,
            "ai_raw_response": response
        }
        return result
    except Exception as e:
        error_result = {
            "success": False,
            "error": f"AI响应处理失败: {str(e)}",
            "ai_raw_response": response
        }
        return error_result

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) 