# 安装依赖：pip install langgraph langchain-openai

from typing import TypedDict, Annotated, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)


# 定义共享状态
class MultiAgentState(TypedDict):
    task: str               # 原始任务
    research_result: str    # 调研 Agent 的输出
    report_draft: str       # 写作 Agent 的输出
    final_output: str       # 格式化 Agent 的输出（最终产物）


# Worker Agent 1：调研 Agent
def research_agent(state: MultiAgentState) -> Command:
    """专注于信息调研，不做写作"""
    prompt = f"""你是一个专业的市场调研分析师。
任务：{state['task']}

请进行调研分析，提供：
1. 核心功能特性对比
2. 定价策略对比
3. 市场定位分析
4. 主要优缺点总结

输出结构化的调研结果，清晰标注各个维度。"""

    result = llm.invoke([HumanMessage(content=prompt)])

    # 更新自己负责的字段，然后把控制权交还给 Supervisor
    return Command(
        update={"research_result": result.content},
        goto="supervisor"
    )


# Worker Agent 2：写作 Agent
def writer_agent(state: MultiAgentState) -> Command:
    """专注于报告写作，不做调研"""
    prompt = f"""你是一个专业的商业分析报告撰写者。

调研结果：
{state['research_result']}

原始任务：{state['task']}

请基于调研结果撰写一份专业的分析报告，包括：
- 执行摘要（核心结论，不超过 200 字）
- 竞品对比分析（用表格）
- 市场机会分析
- 行动建议

报告要逻辑清晰，结论明确。"""

    result = llm.invoke([HumanMessage(content=prompt)])

    return Command(
        update={"report_draft": result.content},
        goto="supervisor"
    )


# Worker Agent 3：格式化 Agent
def formatter_agent(state: MultiAgentState) -> Command:
    """专注于将报告转化为 PPT 提纲，不做内容创作"""
    prompt = f"""你是一个专业的演示文稿设计师。

分析报告：
{state['report_draft']}

请将报告整理为 PPT 提纲：
- 每页幻灯片的标题
- 每页的 3-5 个要点（简洁，适合演讲）
- 建议配图类型

PPT 在 10 页以内，重点突出，适合高管汇报。"""

    result = llm.invoke([HumanMessage(content=prompt)])

    return Command(
        update={"final_output": result.content},
        goto="supervisor"
    )


# Supervisor Agent：任务协调和路由
def supervisor(state: MultiAgentState) -> Command:
    """根据当前状态决定下一步，不亲自执行任务

    这是 Supervisor 模式的核心：通过检查 State 字段是否已填充
    来判断当前进度，决定下一步路由到哪个 Agent
    """
    if not state.get("research_result"):
        # 还没有调研结果，先调研
        print("[Supervisor] 启动调研 Agent...")
        return Command(goto="research_agent")

    elif not state.get("report_draft"):
        # 有调研结果了，写报告
        print("[Supervisor] 调研完成，启动写作 Agent...")
        return Command(goto="writer_agent")

    elif not state.get("final_output"):
        # 有报告了，生成 PPT 提纲
        print("[Supervisor] 报告完成，启动格式化 Agent...")
        return Command(goto="formatter_agent")

    else:
        # 所有步骤完成
        print("[Supervisor] 所有任务完成！")
        return Command(goto=END)


# 构建 Multi-Agent 图
builder = StateGraph(MultiAgentState)

builder.add_node("supervisor", supervisor)
builder.add_node("research_agent", research_agent)
builder.add_node("writer_agent", writer_agent)
builder.add_node("formatter_agent", formatter_agent)

# 所有流程都从 Supervisor 开始，Supervisor 决定路由
builder.add_edge(START, "supervisor")

# 编译
graph = builder.compile()


def run_analysis(task: str) -> dict:
    """执行分析任务"""
    initial_state = MultiAgentState(
        task=task,
        research_result="",
        report_draft="",
        final_output=""
    )

    result = graph.invoke(initial_state)

    return {
        "research": result["research_result"][:500] + "...",  # 截断展示
        "report": result["report_draft"][:500] + "...",
        "ppt_outline": result["final_output"]
    }


# 使用示例
# result = run_analysis(
#     "分析 Notion、Obsidian、Roam Research 三款知识管理工具，分析功能、定价、市场定位"
# )
# print(result["ppt_outline"])