import json
from typing import Dict, Any, Optional

import jingrow
from jingrow.utils.jingrow_cloud import call_ai_model
from jingrow.utils.jinja import render_template

def execute(context=None, inputs=None, config=None):
    """
    AI内容生成节点 - 适配自 jingrow 系统
    执行AI内容生成任务
    """
    try:
        # 参数处理
        if context is None:
            context = {}
        if inputs is None:
            inputs = {}
        if config is None:
            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. 检查模板配置
        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()
        
        # 根据flow_id获取完整的agent对象（通过Jingrow API读取AI Agent记录）
        if not (context and context.get("flow_id")):
            return {
                "success": False,
                "error": "缺少flow_id，无法读取AI Agent记录"
            }
        agent_data = {}
        flow_id = context["flow_id"]
        agent_data = jingrow.get_pg("Local Ai Agent", str(flow_id))
        if not agent_data:
            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. 渲染模板（使用统一的 render_template，容错处理undefined变量）
        try:
            prompt = render_template(template, template_context)
        except Exception as e:
            return {
                "success": False,
                "error": f"模板渲染失败: {str(e)}"
            }

        # 5. 调用AI（模拟AI调用）
        system_message = config.get("ai_system_message", "")
        text_model = config.get("text_model", "gpt-3.5-turbo")
        max_tokens = config.get("max_tokens", 1000)
        temperature = config.get("temperature", 0.9)
        top_p = config.get("top_p", 0.9)

        api_res = 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,
        )
        if not api_res.get("success"):
            return {"success": False, "error": api_res.get("error", "AI模型调用失败")}
        response = api_res.get("response", "")

        # 6. 返回AI输出（对齐jingrow原版格式）
        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

    except Exception as e:
        return {
            "success": False,
            "error": f"AI内容生成节点执行失败: {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)
