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

# ---- 定义系统状态 ----

class MultiAgentState(TypedDict):
    """
    共享状态：所有 Agent 都可以读写这个状态对象。
    每个字段存储一类信息，Agent 通过读写对应字段来协作。
    """
    task: str                         # 原始任务描述
    research_result: str              # 研究员的调研结果
    draft_content: str                # 写作者的初稿
    review_feedback: str              # 审核者的反馈意见
    final_output: str                 # 最终输出
    next_agent: str                   # Supervisor 决定的下一个执行者
    iteration_count: int              # 迭代次数（防止无限循环）
    messages: Annotated[list, operator.add]  # 完整消息历史

MAX_ITERATIONS = 10  # 最大迭代次数，超过则强制结束

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

# ---- 定义三个 Worker Agent ----

def researcher_agent(state: MultiAgentState) -> MultiAgentState:
    """
    研究员 Agent：负责收集和整理与任务相关的背景信息。
    输入：任务描述（state["task"]）
    输出：调研结果（写入 state["research_result"]）
    """
    response = llm.invoke([
        SystemMessage(content=(
            "你是一名专业研究员。请针对给定的任务，收集相关背景信息、"
            "关键数据和重要观点。输出结构化的调研摘要，300字以内。"
        )),
        HumanMessage(content=f"任务：{state['task']}"),
    ])
    print(f"[研究员] 完成调研，字数：{len(response.content)}")
    return {
        "research_result": response.content,
        "messages": [response],
    }


def writer_agent(state: MultiAgentState) -> MultiAgentState:
    """
    写作者 Agent：基于调研结果撰写内容初稿。
    输入：任务描述 + 调研结果
    输出：内容初稿（写入 state["draft_content"]）
    """
    # 如果有审核反馈，则参考反馈修改；否则基于调研结果写初稿
    feedback_section = ""
    if state.get("review_feedback"):
        feedback_section = f"\n\n审核反馈（请据此修改）：\n{state['review_feedback']}"

    response = llm.invoke([
        SystemMessage(content=(
            "你是一名专业内容写作者。请基于提供的调研结果，撰写高质量的文章内容。"
            "要求：逻辑清晰、语言流畅、有实质内容，500字以内。"
        )),
        HumanMessage(content=(
            f"任务：{state['task']}\n\n"
            f"调研结果：{state['research_result']}"
            f"{feedback_section}"
        )),
    ])
    print(f"[写作者] 完成初稿，第 {state['iteration_count']+1} 次写作")
    return {
        "draft_content": response.content,
        "messages": [response],
    }


def reviewer_agent(state: MultiAgentState) -> MultiAgentState:
    """
    审核者 Agent：审查初稿质量，给出通过/修改意见。
    输入：任务描述 + 初稿内容
    输出：审核反馈（写入 state["review_feedback"]）
    特殊约定：如果质量通过，反馈开头写 "APPROVED"
    """
    response = llm.invoke([
        SystemMessage(content=(
            "你是一名严格的内容审核者。请评估文章初稿的质量。\n"
            "评估标准：内容准确性、逻辑连贯性、实用价值。\n"
            "如果质量合格，回复以 'APPROVED' 开头，后接简短评语。\n"
            "如果需要修改，直接说明具体的修改建议，不要写 APPROVED。"
        )),
        HumanMessage(content=(
            f"任务要求：{state['task']}\n\n"
            f"文章初稿：{state['draft_content']}"
        )),
    ])
    print(f"[审核者] 完成审核：{'通过' if 'APPROVED' in response.content else '需修改'}")
    return {
        "review_feedback": response.content,
        "messages": [response],
    }


# ---- 定义 Supervisor Agent ----

def supervisor_agent(state: MultiAgentState) -> MultiAgentState:
    """
    Supervisor：根据当前状态决定下一步派谁去工作。
    路由逻辑：
      - 没有调研结果 → 派研究员
      - 有调研结果但没有初稿 → 派写作者
      - 有初稿但没有审核反馈 → 派审核者
      - 审核通过（APPROVED） → 结束
      - 审核不通过 → 派写作者修改（检查迭代次数防止死循环）
    """
    # 超过最大迭代次数，强制结束（防止无限循环）
    if state["iteration_count"] >= MAX_ITERATIONS:
        print(f"[Supervisor] 已达最大迭代次数 {MAX_ITERATIONS}，强制结束")
        return {
            "next_agent": "end",
            "final_output": state.get("draft_content", "任务超时，未能完成"),
        }

    # 路由决策逻辑
    if not state.get("research_result"):
        print("[Supervisor] → 派遣研究员")
        next_step = "researcher"

    elif not state.get("draft_content"):
        print("[Supervisor] → 派遣写作者")
        next_step = "writer"

    elif not state.get("review_feedback"):
        print("[Supervisor] → 派遣审核者")
        next_step = "reviewer"

    elif "APPROVED" in state.get("review_feedback", ""):
        print("[Supervisor] → 内容已通过审核，任务完成")
        next_step = "end"

    else:
        # 审核未通过，让写作者根据反馈修改
        print(f"[Supervisor] → 审核未通过，第 {state['iteration_count']+1} 次修改")
        # 清空旧的审核反馈，让写作者重写后再次审核
        next_step = "writer"

    return {
        "next_agent": next_step,
        "iteration_count": state["iteration_count"] + 1,
    }


# ---- 路由函数：Supervisor 的决策映射到图的边 ----

def route_by_supervisor(state: MultiAgentState) -> Literal[
    "researcher", "writer", "reviewer", "__end__"
]:
    """
    根据 Supervisor 的决策，返回下一个要执行的节点名称。
    LangGraph 用这个函数的返回值决定走哪条边。
    """
    next_agent = state.get("next_agent", "researcher")
    if next_agent == "end":
        return "__end__"
    return next_agent  # 返回 "researcher"、"writer" 或 "reviewer"


# ---- 构建 LangGraph 图 ----

def build_supervisor_graph():
    builder = StateGraph(MultiAgentState)

    # 注册所有节点
    builder.add_node("supervisor", supervisor_agent)
    builder.add_node("researcher", researcher_agent)
    builder.add_node("writer", writer_agent)
    builder.add_node("reviewer", reviewer_agent)

    # 入口：从 Supervisor 开始
    builder.add_edge(START, "supervisor")

    # Supervisor 根据路由函数决定下一步
    builder.add_conditional_edges(
        "supervisor",
        route_by_supervisor,
        {
            "researcher": "researcher",
            "writer": "writer",
            "reviewer": "reviewer",
            "__end__": END,
        }
    )

    # 所有 Worker 完成后，回到 Supervisor 汇报
    builder.add_edge("researcher", "supervisor")
    builder.add_edge("writer", "supervisor")
    builder.add_edge("reviewer", "supervisor")

    return builder.compile()


# ---- 运行示例 ----

def run_multi_agent_task(task: str) -> str:
    graph = build_supervisor_graph()

    initial_state = {
        "task": task,
        "research_result": "",
        "draft_content": "",
        "review_feedback": "",
        "final_output": "",
        "next_agent": "",
        "iteration_count": 0,
        "messages": [],
    }

    print(f"=== 开始多 Agent 协作任务 ===\n任务：{task}\n")

    final_state = graph.invoke(initial_state)

    print("\n=== 任务完成 ===")
    print(f"总迭代次数：{final_state['iteration_count']}")
    print(f"\n最终内容：\n{final_state['draft_content']}")

    return final_state["draft_content"]


# 测试
result = run_multi_agent_task("写一篇面向Java程序员的AI Agent入门介绍，500字")

# 好的设计：双保险
if task_is_complete(state):      # 正常结束：任务完成
    return END
if state["iterations"] > 20:    # 兜底：超过最大次数强制结束
    return END

# 坏的设计：审核者直接修改初稿
def reviewer_agent(state):
    state["draft_content"] = improved_draft  # 不应该！审核者应该写反馈，不应该改稿

# 好的设计：审核者只写反馈
def reviewer_agent(state):
    return {"review_feedback": "建议增加具体数据支撑..."}  # 只修改自己负责的字段

def supervisor_with_error_handling(state):
    try:
        # 正常路由逻辑
        return route_to_next_agent(state)
    except AgentTimeoutError as e:
        # Worker 超时：记录日志，尝试换一个 Worker
        log_error(e)
        return {"next_agent": "fallback_agent"}
    except CriticalError as e:
        # 严重错误：终止任务，返回错误信息给用户
        return {"next_agent": "end", "final_output": f"任务失败：{e}"}

import subprocess
import tempfile
import os

# 代码生成 Agent：根据需求写代码
def code_generator(state):
    feedback = state.get("review_comments", "")
    prompt = f"需求：{state['requirement']}"
    if feedback:
        prompt += f"\n\n代码审查意见（请据此修改）：{feedback}"

    response = llm.invoke([
        SystemMessage(content=(
            "你是一名 Python 专家。请编写符合需求的 Python 代码。"
            "只输出代码，不要解释，不要 markdown 代码块格式。"
        )),
        HumanMessage(content=prompt),
    ])
    print("[代码生成] 完成代码编写")
    return {"generated_code": response.content}


# 测试 Agent：为生成的代码编写单元测试并执行
def test_agent(state):
    # 请 LLM 为代码生成测试用例
    response = llm.invoke([
        SystemMessage(content=(
            "你是测试工程师。请为以下代码编写 pytest 单元测试。"
            "只输出测试代码，不要解释。"
        )),
        HumanMessage(content=f"待测代码：\n{state['generated_code']}"),
    ])
    test_code = response.content

    # 实际执行测试（在临时目录中）
    test_result = "PASS"  # 默认通过
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            # 写入代码文件
            code_file = os.path.join(tmpdir, "solution.py")
            test_file = os.path.join(tmpdir, "test_solution.py")
            with open(code_file, "w") as f:
                f.write(state["generated_code"])
            with open(test_file, "w") as f:
                f.write(test_code)

            # 运行 pytest
            result = subprocess.run(
                ["python", "-m", "pytest", test_file, "-v", "--tb=short"],
                capture_output=True, text=True, timeout=30, cwd=tmpdir
            )
            test_result = "PASS" if result.returncode == 0 else f"FAIL:\n{result.stdout}"
    except Exception as e:
        test_result = f"ERROR: {e}"

    print(f"[测试] 运行结果：{test_result[:50]}")
    return {"test_result": test_result, "test_code": test_code}


# 代码审查 Agent：检查代码质量
def code_reviewer(state):
    response = llm.invoke([
        SystemMessage(content=(
            "你是高级代码审查专家。请审查代码质量：\n"
            "检查项：命名规范、错误处理、边界条件、代码简洁性。\n"
            "如果代码质量合格且测试通过，回复以 'LGTM'（Looks Good To Me）开头。\n"
            "否则，直接给出具体修改意见。"
        )),
        HumanMessage(content=(
            f"代码：\n{state['generated_code']}\n\n"
            f"测试结果：{state.get('test_result', '未测试')}"
        )),
    ])

    approved = response.content.startswith("LGTM")
    print(f"[代码审查] {'通过' if approved else '未通过'}")
    return {"review_comments": response.content, "code_approved": approved}