import json
import os

from alayaflow.api import Flow
from dotenv import load_dotenv
load_dotenv()

def main():
    # Initialize flow api singleton

    flow = Flow()
    flow.init({
        "alayahub_url": "http://your-alayahub-url",
    })

    chat_model_api_key = os.getenv("DEEPSEEK_API_KEY")
    if not chat_model_api_key:
        raise ValueError("请设置环境变量 DEEPSEEK_API_KEY")

    extract_model_api_key = os.getenv("TENCENT_DEEPSEEK_API_KEY")
    if not extract_model_api_key:
        raise ValueError("请设置环境变量 TENCENT_DEEPSEEK_API_KEY")

    jina_api_key = os.getenv("JINA_API_KEY")
    if not jina_api_key:
        raise ValueError("请设置环境变量 JINA_API_KEY")
    
    flow.register_models([
        {
            # Local used fields
            "name": "DeepSeek Chat",
            "model_id": "deepseek-chat",
            "provider_name": "DeepSeek",
            # Connection credentials
            "model_name": "deepseek-chat",
            "base_url": "https://api.deepseek.com/v1",
            "api_key": chat_model_api_key,
        },
        {
            # Local used fields
            "name": "DeepSeek Extract",
            "model_id": "deepseek-v3",
            "provider_name": "DeepSeek",
            # Connection credentials
            "model_name": "deepseek-v3.1-terminus",
            "base_url": "https://api.lkeap.cloud.tencent.com/v1",
            "api_key": extract_model_api_key,
        },
        {
            "name": "Jina Rerank",
            "model_id": "jina-reranker",
            "provider_name": "Jina",
            "model_type": "JinaRerank",
            "model_name": "jina-reranker-v3",
            "base_url": "null",
            "api_key": jina_api_key,
        }
    ])

    workflow_id = 'autotable'
    workflow_version = '1.0.0'
    init_args = {
        "alayamem_url": "http://100.64.0.5:5556",
    }
    flow.load_workflow(workflow_id, workflow_version, init_args)

    # 新的输入格式：使用 fields 字段，支持递归结构
    input_data = {
        "fields": [
            {
                "申请人信息": [
                "姓名",
                "性别",
                "出生年月",
                "民族",
                "学位",
                "职称",
                "是否在站博士后",
                "国别或地区",
                "申请人类别",
                "工作单位",
                "主要研究领域",
                {
                    "联系方式": [
                    "电子邮箱",
                    "办公电话"
                    ]
                }
                ]
            },
            {
                "依托单位信息": [
                "名称",
                "联系人",
                "网站地址",
                {
                    "联系方式": [
                    "电子邮箱",
                    "电话"
                    ]
                }
                ]
            },
            "合作研究单位信息 (单位名称)",
            {
                "项目基本信息": [
                "项目名称",
                "英文名称",
                "资助类别",
                "亚类说明",
                "附注说明",
                "申请代码",
                "研究期限",
                "研究方向",
                "申请资助经费",
                "研究属性",
                "中文关键词",
                "英文关键词"
                ]
            },
            {
                "摘要": [
                "中文摘要",
                "英文摘要"
                ]
            }
        ],
    }
    context = {
        "user_id": "user",
        "profile_id": "pf_c1186db4f8fc4c998a69d31a236718f9",
        "extract_model_id": "deepseek-v3"
    }
        
    final_result = None
    for event in flow.exec_workflow(workflow_id, workflow_version, input_data, context):
        if "error" in event:
            print(f"\n[Error] {event['error']}")
            break
        # 查找包含最终状态的事件
        # LangGraph 的 astream_events 可能在 data.output 或 data 中包含最终状态
        if "data" in event:
            data = event["data"]
            # 检查 output 字段（通常是最终状态）
            if isinstance(data, dict) and "output" in data:
                output = data["output"]
                if isinstance(output, dict) and ("context_by_task" in output or "final_result" in output):
                    final_result = output
            # 或者直接检查 data 中是否包含最终状态字段
            elif isinstance(data, dict) and ("context_by_task" in data or "final_result" in data):
                final_result = data
    
    # 打印最终结果
    if final_result:
        print("\n最终结果：")
        # print(json.dumps(final_result, indent=2, ensure_ascii=False))
        print(final_result.get("final_result", {}))
        
        # 打印错误信息（如果有）
        errors = final_result.get("errors", {})
        if errors:
            print("\n错误信息：")
            print(json.dumps(errors, indent=2, ensure_ascii=False))

if __name__ == '__main__':
    main()