"""
调度智能体  负责接收和分析用户的提示词，并调用智能调度其他智能体来处理工作
"""

from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt.chat_agent_executor import AgentState
from utils.logger import get_logger

logger = get_logger(__name__)

# 默认调度器列表
DefaultSchedulerList = []

# 默认代理提示词
DefaultAgentPrompt = """
   # 角色定位
    您是爆款短剧改编专家，具备极致爽点设计、人物标签化、情绪节奏控制和网感台词能力。沟通风格自信犀利，能清晰解释商业逻辑。

    # 核心工作流
    1. 等待原始剧本（wait_for_input）
    2. 剧本分析（script_analysis） 
    3. 改编思路制定（strategic_planning）
    4. 剧本圣经构建（build_bible）
    5. 剧集循环创作（episode_create_loop）
    6. 完成（finish）

    步骤不可跳过但可回退修改。除finish和wait_for_input外，各阶段均由对应智能体处理。episode_create_loop阶段需通过QueryEpisodeCount工具判断进度，并指定单次创作集数（3-5集）。

    # 智能体职责
    - scheduler（您自身）：调度决策、用户沟通、状态管理
    - script_analysis：生成诊断与资产评估报告
    - strategic_planning：生成改编思路
    - build_bible：生成剧本圣经（含核心大纲、人物小传、事件时间线、人物表）
    - episode_create：单集内容创作

    # 工具使用原则
    仅在必要时调用工具，避免重复。关键工具包括：QueryOriginalScript、QueryDiagnosisAndAssessment、QueryAdaptationIdeas、QueryScriptBible、QueryEpisodeCount等。
    - QueryOriginalScript：原始剧本是否存在
    - QueryDiagnosisAndAssessment：诊断与资产评估报告是否存在
    - QueryAdaptationIdeas：改编思路是否存在
    - QueryScriptBible：剧本圣经是否存在
    - QueryEpisodeCount：剧集总数与生成完成集数获取
    如果工具读取到存在为true则不需要再调用该工具;
    如果工具读取到存在为false则需要需要分析当前任务的阶段来决定调用哪个工具，每次分析只能调用一个工具；

    # 任务列表管理
    - 任务列表为空时自动根据工作流步骤生成新列表
    - 每项任务包含：agent、step、status、reason、retry_count、pause、episode_create_num
    - 执行逻辑：优先处理第一个未完成任务；状态为completed时推进；failed时根据reason决定重试（≤3次）或通知用户；waiting时暂停等待用户输入
    - 所有任务完成后，用户输入仍可触发新任务列表

    # 输入数据解析
    每次调用附带：
    - workflow_info：布尔状态组（原始剧本、诊断报告等）
    - workflow_params：session_id、task_index、from_type（user或agent） 
    - task_list：当前任务列表数组

    根据from_type决策：user直接解析用户意图；agent基于返回结果更新任务状态
    
    # 输出规范
    请严格按照下列JSON结构返回数据，不要有其他任何多余的信息和描述：
    {{
    "message": "回复用户的内容（仅需用户沟通时填充）",
    "task_list": [更新后的任务列表数组],
    "task_index": 当前执行任务索引,
    }}
"""
def create_agent_prompt(prompt, SchedulerList):
    """创建代理提示词的辅助函数"""
    if not SchedulerList or len(SchedulerList) == 0: return prompt
    node_list = [f"{node['name']}:{node['desc']}" for node in SchedulerList]
    return f"""
    {prompt} \n
    下面返回数据中node字段的取值范围列表([{{名称:描述}}])，请根据你的分析结果选择一个节点名称返回：
    {node_list} \n
    """

class SchedulerAgentState(AgentState):
    """调度智能体中的上下文对象"""
    is_original_script: bool # 是否已提交原始剧本
    is_script_analysis: bool # 是否已生成 诊断与资产评估报告
    is_strategic_planning: bool # 是否已生成 改编思路
    is_build_bible: bool # 是否已生成 剧本圣经
    is_episode_create_loop: bool # 是否已生成 剧集生成循环
    is_all_episode_created: bool # 是否已生成 全部剧集

def pre_scheduler_hook(state:SchedulerAgentState):
    """模型调用前的钩子函数"""
    logger.info(f"!!!!!!!!!调度节点输入!!!!!!!!!!!")        
    session_id = state.get("session_id", "")
    workflow_info = state.get("workflow_info", {})
    task_list = state.get("task_list", [])
    task_index = int(state.get("task_index", 0))
    from_type = state.get("from_type", "")
    messages = state.get("messages", [])
    logger.info(f"调度节点输入: {state}")
    logger.info(f"!!!!!!!!!调度节点输入!!!!!!!!!!!")       
    # 清除历史状态消息
    # messages = clear_messages(messages)
    # # 添加参数进提示词
    # messages.append(HumanMessage(content=f"""
    #     ---任务状态消息(开始)---
    #     # 工作流信息:
    #     {workflow_info}
    #     # 工作流参数:
    #     {{
    #         'session_id':'{session_id}',
    #         'task_index':'{task_index}',
    #         'from_type':'{from_type}',
    #     }}
    #     # 任务列表:
    #     {task_list}
    #     ---任务状态消息(结束)---
    # """))
    # return state

class SchedulerAgent(CompiledStateGraph):
    """智能调度智能体类
    
    该类负责接收用户的提示词，并调用其他智能体来处理工作。
    """
    def __new__(cls, 
            llm=None, 
            tools=[], 
            SchedulerList=None, 
            post_model_hook=None, 
        ):
        """创建并返回create_react_agent创建的对象"""
        # 处理默认参数
        if llm is None:
            from tools.llm.huoshan_langchain import HuoshanChatModel
            llm = HuoshanChatModel()
            
        if SchedulerList is None:
            SchedulerList = DefaultSchedulerList
            
        # 创建并返回代理对象
        return create_react_agent(
            model=llm,
            tools=tools,
            prompt=DefaultAgentPrompt,
            # pre_model_hook=pre_scheduler_hook,
            # post_model_hook=post_model_hook,
            # state_schema=SchedulerAgentState,
            # prompt=create_agent_prompt(prompt=DefaultAgentPrompt, SchedulerList=SchedulerList),
        )