import os
import operator
from typing import Annotated, List, Tuple, TypedDict, Union, Literal

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain.agents import create_agent
from pydantic import BaseModel, Field

#https://tavily.com/
#uv pip install tavily-python -i https://pypi.tuna.tsinghua.edu.cn/simple

# ---- 配置区：请替换为你自己的 Key ----
os.environ["TAVILY_API_KEY"] = "your-tavily-key-here"
DASHSCOPE_API_KEY = "your-dashscope-key-here"

# 初始化搜索工具，max_results 控制每次搜索返回的条数
search_tool = TavilySearchResults(max_results=3)
tools = [search_tool]

# 规划阶段用推理能力强的 Qwen-Max
planner_llm = ChatTongyi(
    model="qwen-max",
    temperature=0.1,          # 低温度保证规划稳定
    api_key=DASHSCOPE_API_KEY,
)

# 执行阶段可以用更经济的模型，这里用 Qwen-Plus
executor_llm = ChatTongyi(
    model="qwen-plus",
    temperature=0.01,
    api_key=DASHSCOPE_API_KEY,
)

# 创建执行 Agent—— 它就是一个标准的 ReAct Agent
executor_agent = create_agent(
    model=executor_llm,
    tools=tools,
    system_prompt="你是一个执行力很强的助手。请根据给定的任务步骤，准确完成并给出结果。",
)


class PlanExecute(TypedDict):
    """贯穿整个工作流的状态对象"""
    input: str                                              # 用户原始输入
    plan: List[str]                                         # 当前计划步骤列表
    # past_steps 用 operator.add 做累加——每次节点返回新步骤时自动追加到历史
    past_steps: Annotated[List[Tuple[str, str]], operator.add]
    response: str                                           # 最终回复（非空时流程结束）


class Plan(BaseModel):
    """规划器输出的结构化计划"""
    steps: List[str] = Field(
        description="需要顺序执行的步骤列表，每一步必须独立且包含完整上下文"
    )


class Response(BaseModel):
    """直接回复用户的内容"""
    response: str


class Act(BaseModel):
    """
    重规划器的决策输出——要么回复用户，要么给出更新后的计划。
    用 Union 类型确保一次只走一条分支。
    """
    action: Union[Response, Plan] = Field(
        description="如果任务已完成，返回 Response；如果还需要继续，返回 Plan"
    )



# 规划器的系统提示——核心是"拆得合理、不越界"
planner_prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        """你是一个擅长任务分解的规划专家。对于用户提出的目标，请将其拆解为顺序执行的步骤列表。

规则：
1. 每一步必须自包含——执行者不需要上下文就能理解这步要做什么
2. 不要添加无关步骤，最后一步的结果应该直接导向最终答案
3. 优先使用中文描述步骤
4. 步骤数量控制在 3~6 步，过少说明拆分不到位，过多说明混入了无关操作""",
    ),
    ("placeholder", "{messages}"),
])

# 将提示词模板和 LLM 拼成一条链，with_structured_output 让 LLM 直接输出 Plan 对象
planner_chain = planner_prompt | planner_llm.with_structured_output(Plan)


def plan_step(state: PlanExecute) -> dict:
    """规划节点：分析用户原始输入，生成执行计划"""
    result = planner_chain.invoke({"messages": [("user", state["input"])]})
    return {"plan": result.steps}


result = planner_chain.invoke({
    "messages": [(
        "user",
        "帮我调研一下2025年国内主流代码助手工具（通义灵码、文心快码、CodeGeeX），"
        "从代码补全准确率、支持语言、价格三个维度做对比，最后给出选型建议。"
    )]
})

for i, step in enumerate(result.steps, 1):
    print(f"  Step {i}: {step}")


def execute_step(state: PlanExecute) -> dict:
    """
    从当前计划中取出第一步，交给执行 Agent 去跑。
    执行完后把 (步骤描述, 执行结果) 追加到 past_steps。
    """
    plan = state["plan"]

    # 把完整计划格式化，让 Agent 知道"我在做什么、后面还有什么"
    plan_overview = "\n".join(f"  {i+1}. {step}" for i, step in enumerate(plan))
    current_task = plan[0]  # 只取第一步

    # 拼装给 Agent 的指令：先展示全貌，再指定当前任务
    task_prompt = f"""以下是完整的执行计划：
{plan_overview}

现在请你只执行第 1 步，不要做后续步骤：
  → {current_task}

请直接给出这一步的执行结果，不要啰嗦。"""

    agent_result = executor_agent.invoke({
        "messages": [("user", task_prompt)]
    })

    # 取 Agent 最后一条消息的内容作为这一步的产出
    step_output = agent_result["messages"][-1].content

    return {
        # 这一步走完后，计划需要弹出已执行的第一项——这个逻辑在 replan_step 里
        "past_steps": [(current_task, step_output)],
    }


replanner_prompt = ChatPromptTemplate.from_template("""
你是一个项目进度管控专家。根据以下信息，判断当前任务的状态并做出决策。

## 用户原始目标
{input}

## 原始计划
{plan}

## 已完成的步骤及结果
{past_steps}

## 决策规则
1. 如果所有步骤已完成且结果充分，直接给出最终答案（使用 Response）
2. 如果还有步骤未执行，返回更新后的计划（使用 Plan），注意：
   - 只保留尚未执行的步骤
   - 如果已完成步骤的结果表明原计划的后续步骤需要调整，直接修改
   - 不要返回已经做过的步骤
3. 如果某个步骤执行失败了，请在计划中补充重试或替代方案
""")

replanner_chain = replanner_prompt | planner_llm.with_structured_output(Act)


def replan_step(state: PlanExecute) -> dict:
    """审视执行进度，决定下一步走向"""
    output = replanner_chain.invoke(state)

    if isinstance(output.action, Response):
        # 任务完成，直接返回给用户
        return {"response": output.action.response}
    else:
        # 还有步骤要执行，更新计划列表
        return {"plan": output.action.steps}
    


from langgraph.graph import StateGraph, START, END

# ---- 第一步：创建状态图，绑定状态类型 ----
workflow = StateGraph(PlanExecute)

# ---- 第二步：注册三个核心节点 ----
workflow.add_node("planner", plan_step)       # 规划节点
workflow.add_node("executor", execute_step)   # 执行节点
workflow.add_node("replanner", replan_step)   # 重规划节点

# ---- 第三步：编排边的流向 ----
workflow.add_edge(START, "planner")           # 入口 → 规划
workflow.add_edge("planner", "executor")      # 规划 → 执行
workflow.add_edge("executor", "replanner")    # 执行 → 重规划

# ---- 第四步：条件分支——重规划后走哪条路 ----
def should_continue(state: PlanExecute) -> Literal["executor", "__end__"]:
    """如果 response 字段非空，说明重规划器认为任务完成，结束流程"""
    if state.get("response"):
        return END
    return "executor"

workflow.add_conditional_edges("replanner", should_continue, {
    "executor": "executor",
    "__end__": END,
})



inputs = {
    "input": (
        "帮我调研2025年国内三大代码助手工具（通义灵码、文心快码、CodeGeeX），"
        "从代码补全准确率、支持语言、价格三个维度做对比，最后给出选型建议。"
        "请用中文输出最终结果。"
    )
}

# stream 模式可以看到每一步的中间输出
for event in app.stream(inputs):
    for node_name, node_output in event.items():
        print(f"\n{'='*50}")
        print(f"📍 当前节点: {node_name}")
        # 避免打印太长的内容
        for key, value in node_output.items():
            if isinstance(value, str) and len(value) > 200:
                print(f"  {key}: {value[:200]}...")
            else:
                print(f"  {key}: {value}")