"""Agent 节点：核心推理与工具调用"""

from typing import Dict, Any, Optional
from langchain_core.messages import SystemMessage, AIMessage, AIMessageChunk
from langchain_core.runnables.config import RunnableConfig
from ..state import AgentState
from backend.app.logger import info, warning, error
from .stream_context import token_queue_var


# 系统提示词（从 main_graph_builder.py 搬过来）
SYSTEM_PROMPT = """你是一个智能助手，可以使用多种工具完成复杂任务。你必须用中文回复。

## 核心工具与能力
你可以使用以下工具（函数），但只能在真正需要时调用，禁止无意义的测试调用或重复调用：
1. rag_search – 从内部知识库中检索文档，输入为优化后的查询字符串。
2. web_search – 联网搜索获取最新信息，输入为搜索关键词。
3. contact_lookup – 查询企业通讯录，输入姓名、部门或邮箱等。
4. dictionary_lookup – 翻译单词、查询词典或提取术语。
5. news_analysis – 获取或分析新闻资讯。

## 工作流程（ReAct 决策闭环）
你必须严格按照思考 → 行动 → 观察的闭环来处理每个请求，具体规则如下：

### 1. 初始决策
- 如果用户的问题很明确且你已有足够内部知识，可以直接回答，无需调用任何工具。
- 如果需要外部信息，请按以下优先级选择工具：
  - 优先使用 rag_search。
  - 若第一次 rag_search 返回的结果不相关或质量低，你可以改写查询关键词再次调用 rag_search（最多重复一次）。
  - 如果两次 rag_search 均无法获得满意信息，或者用户明确要求实时资讯，则必须切换为 web_search。
- 遇到通讯录、词典、新闻类明确需求，直接调用对应的专用工具。

### 2. 观察与反思
- 每次工具调用返回结果后，你必须先评估结果质量（内容是否相关、是否充分）。
- 如果信息不足，根据上述规则决定下一步行动；如果信息足够，则直接生成最终答案，绝不再调用任何工具。
- 在整个过程中，禁止使用工具返回的信息直接重复或编造来源，必须如实标注。

### 3. 结束条件
当你认为已经拥有足够信息回答用户时，输出最终回复并停止调用工具。若连续调用工具超过 5 轮仍未解决，也必须基于当前收集到的信息给出最佳回答并说明局限性。

## 回答规范
1. 来源标注：回答开头用方括号注明信息来源，如多处来源按使用顺序列出：
   - 知识库：【知识库：相关文档主题】
   - 联网搜索：【联网搜索：来源网站或摘要】
2. 思维链：对于需要复杂推理的问题，请将推理过程放在 <think>...</think> 标签内，并置于回答最前面（来源标注之前）。
3. 内容要求：回答应重点突出、条理清晰，优先结合用户背景信息进行个性化；若无任何可靠依据，如实说明“暂时无法回答”。

## 特别注意
- 不要向用户暴露任何工具调用的技术细节（如参数、函数名）。
- 如果用户只是闲聊、问候或道别，直接友好回复，严禁调用任何工具。
- 所有联网搜索必须以获取帮助用户为目的，不得搜索无关内容。

现在，请遵循以上规则处理用户的每一次输入。记住：思考 → 行动 → 观察 → 直到完成。"""


def create_agent_node(llm_with_tools, llm):
    """创建 Agent 节点函数"""

    async def agent_node(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
        """
        Agent 节点：调用带工具的 LLM，处理步数限制

        Args:
            state: 当前状态
            config: 运行配置

        Returns:
            状态更新字典
        """
        current_step = state.get("current_step", 0)
        info(f"[Agent] 第 {current_step} 步推理")

        try:
            # 组装完整消息：系统提示 + 历史消息
            full_messages = [SystemMessage(content=SYSTEM_PROMPT)] + state.messages
            
            info(f"[Agent] 消息数量: {len(full_messages)}, 最后一条: {type(full_messages[-1]).__name__}")

            # 判断是否达到步数上限
            if current_step >= state.get("max_steps", 10):
                info(f"[Agent] 达到步数上限 {state.get('max_steps', 10)}，强制结束，不绑定工具")
                current_llm = llm.bind_tools([])
            else:
                current_llm = llm_with_tools

            info(f"[Agent] 调用带工具的 LLM...")

            # 获取 token 队列
            token_queue = token_queue_var.get()
            if token_queue is None:
                error("[Agent] ❌ token_queue 为 None！")
                raise RuntimeError("token_queue 上下文变量未设置")

            # 完整消息
            full_content = ""
            full_reasoning_content = ""
            full_tool_calls = []

            # 流式调用 LLM
            async for chunk in current_llm.astream(full_messages):
                if isinstance(chunk, AIMessageChunk):
                    # 处理 content
                    if chunk.content:
                        full_content += chunk.content
                        await token_queue.put({
                            "type": "llm_token",
                            "node": "agent",
                            "token": chunk.content,
                            "reasoning_token": "",
                            "turn": current_step,
                            "phase": "answering" if not full_tool_calls else "thinking"
                        })

                    # 处理 reasoning_content
                    if hasattr(chunk, 'additional_kwargs') and chunk.additional_kwargs:
                        reasoning_content = chunk.additional_kwargs.get("reasoning_content", "")
                        if reasoning_content:
                            full_reasoning_content += reasoning_content
                            await token_queue.put({
                                "type": "llm_token",
                                "node": "agent",
                                "token": "",
                                "reasoning_token": reasoning_content,
                                "turn": current_step,
                                "phase": "thinking"
                            })

                    # 处理 tool_calls
                    if hasattr(chunk, 'tool_calls') and chunk.tool_calls:
                        # 合并 tool_calls
                        for tc in chunk.tool_calls:
                            # 查找是否已经有这个 id 的 tool_call
                            found = False
                            for existing_tc in full_tool_calls:
                                if existing_tc.get("id") == tc.get("id"):
                                    # 合并 args
                                    existing_tc["args"] = {**existing_tc.get("args", {}), **tc.get("args", {})}
                                    found = True
                                    break
                            if not found:
                                full_tool_calls.append(tc)
                                # 发送工具调用开始事件
                                await token_queue.put({
                                    "type": "tool_call_start",
                                    "tool": tc.get("name"),
                                    "args": tc.get("args"),
                                    "id": tc.get("id", ""),
                                    "turn": current_step
                                })

            # 构建完整的 AIMessage
            response = AIMessage(
                content=full_content,
                tool_calls=full_tool_calls if full_tool_calls else None
            )
            if full_reasoning_content:
                response.additional_kwargs["reasoning_content"] = full_reasoning_content

            info(f"[Agent] LLM 调用成功！响应类型: {type(response).__name__}")
            if hasattr(response, 'tool_calls') and response.tool_calls:
                info(f"[Agent] 检测到工具调用: {[tc['name'] for tc in response.tool_calls]}")

            # 返回状态更新
            return {
                "messages": [response],
                "current_step": current_step + 1,
                "llm_calls": state.get("llm_calls", 0) + 1
            }

        except Exception as e:
            error(f"[Agent] ❌ 第 {current_step} 步推理出错: {e}")
            import traceback
            error(f"[Agent] 堆栈: {traceback.format_exc()}")
            # 发送错误事件
            token_queue = token_queue_var.get()
            if token_queue:
                await token_queue.put({
                    "type": "error",
                    "message": str(e)
                })
            raise

    return agent_node