# ！/usr/bin/python
# -*- coding: utf-8 -*-
'''
@Project    : 
@File       :workflow_node.py
@IDE        :Cursor
@Author     :LINGMIN
@Date       :2025/08/10 18:00
'''


import json
import sys
from foundation.logger.loggering import server_logger
from foundation.utils.common import handler_err
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from foundation.agent.workflow.test_cus_state import TestCusState
from foundation.agent.generate.test_intent import intent_identify_client
from foundation.agent.test_agent import test_agent_client
from foundation.schemas.test_schemas import FormConfig
from foundation.agent.generate.model_generate import generate_model_client
from foundation.utils.yaml_utils import system_prompt_config



class TestWorkflowNode:
    """
        工作流节点定义
    """
    def __init__(self):
        """初始化模型和会话管理"""

    

    def supervisor_agent(self , state: TestCusState):
        """
            每个代理都与一个 Supervisor 代理通信（主管代理）。由  Supervisor 代理决定接下来应调用哪个代理
            :param state:
            :return:
        """
        session_id = state["session_id"]
        trace_id = state["trace_id"]
        user_input = state["user_input"]
        route_next = state.get("route_next")
        
        server_logger.info(trace_id=trace_id, msg=f"\n===================================[Supervisor].begin-route_next:{route_next}=============================")
        
        config = {
            "session_id": session_id
        }
        # 格式化输出，智能格式化输出
        route_next = intent_identify_client.recognize_intent(trace_id=trace_id , config=config , input=user_input)
        server_logger.info(trace_id=trace_id, msg=f"[Supervisor].intent_identify_client.recognize_intent:{route_next}")
        if route_next not in ["chat_box_generate" , "common_agent"]:
            route_next = "chat_box_generate"

        
        server_logger.info(trace_id=trace_id, msg=f"\n===================================[Supervisor].end-route_next:{route_next}=============================")
        return {
            "route_next": route_next
        }



    async def common_agent_node(self , state: TestCusState):
        """
            通用代理节点
            :param state:
            :return:
        """
        session_id = state["session_id"]
        trace_id = state["trace_id"]
        user_input = state["user_input"]
        config_param = FormConfig(session_id=session_id)
        task_prompt_info = {"task_prompt": ""}
        response_content = await test_agent_client.handle_query(trace_id=trace_id , config_param=config_param, 
                                                                task_prompt_info=task_prompt_info, 
                                                                input_query=user_input, context=None)
        messages = [AIMessage(content=response_content, name="common_agent_node")]
        return {
            "messages": messages,
            "previous_agent": "common_agent",
            "route_next": "FINISH"   # ✅ 直接结束流程
        }
    

    async def chat_box_generate(self , state: TestCusState) -> dict:
        """
            模型生成节点（纯生成类问题）
            :param state:
            :return:
        """
        session_id = state["session_id"]
        trace_id = state["trace_id"]
        user_input = state["user_input"]
        task_prompt_info = state["task_prompt_info"]
        task_prompt_info["task_prompt"] = ""

      # 创建ChatPromptTemplate
        template = ChatPromptTemplate.from_messages([
            ("system", system_prompt_config['system_prompt']),
            ("user", user_input)
        ])

        task_prompt_info = {"task_prompt": template}

        response_content = await generate_model_client.get_model_generate_invoke(trace_id=trace_id , task_prompt_info=task_prompt_info)
        messages = [AIMessage(content=response_content , name="chat_box_generate")]
        server_logger.info(trace_id=trace_id, msg=f"【result】: {response_content}", log_type="chat_box_generate")
        return {
            "messages": messages,
            "route_next": "FINISH"   # ✅ 直接结束流程
        }