"""三种 Multi-Agent 模式的真实执行引擎 — 全部真实 LLM 调用"""
import json, time, uuid
from agent import Agent

# ── Agent 定义 ──

SUPERVISOR = Agent(
    name="Supervisor", icon="👔",
    system_prompt=(
        "你是多 Agent 团队的主管。你的职责是拆解任务、分配给专家、整合结果。\n\n"
        "你有两个专家：\n"
        "- Researcher: 信息搜索专家，擅长用 web_search 搜索真实信息\n"
        "- Coder: 数据分析专家，擅长用 calculator 和 run_python 做计算和分析\n\n"
        "工作流程：\n"
        "1. 收到任务后，先分析需要哪些子任务\n"
        "2. 把子任务分配给合适的专家（通过给专家发送清晰的指令）\n"
        "3. 收到专家结果后，整合成高质量的最终答案\n\n"
        "重要：你不要自己搜索或计算，必须委托给专家。你只负责调度和整合。"
    ),
    tool_names=[],
)

RESEARCHER = Agent(
    name="Researcher", icon="🔍",
    system_prompt=(
        "你是信息搜索专家。你的唯一工具是 web_search。\n\n"
        "规则：\n"
        "1. 收到任务后，用 web_search 搜索，只搜索 1-2 次，不要重复搜索\n"
        "2. 第一次搜索用核心关键词，第二次只在第一次结果不够时才搜\n"
        "3. 返回结构化的搜索结果，包含关键数据\n"
        "4. 不要编造数据，搜不到就说搜不到\n"
        "5. 搜索完成后立即返回结果，不要继续搜索"
    ),
    tool_names=["web_search"],
)

CODER = Agent(
    name="Coder", icon="💻",
    system_prompt=(
        "你是数据分析师。你的工具是 calculator 和 run_python。\n\n"
        "规则：\n"
        "1. 收到数据后，用 calculator 做简单计算，用 run_python 做复杂分析\n"
        "2. 返回清晰的计算结果和分析结论\n"
        "3. 数据来自上一步的搜索结果，不要自己搜索\n"
        "4. 如果需要生成报告，用 run_python 格式化输出"
    ),
    tool_names=["calculator", "run_python"],
)

DEBATE_PRO = Agent(
    name="正方辩手", icon="🟢",
    system_prompt=(
        "你是正方辩手。针对给定辩题，你需要搜索资料并提出支持性论点。\n\n"
        "规则：\n"
        "1. 用 web_search 搜索真实案例和数据支持你的论点\n"
        "2. 每个论点要有事实依据，不能空谈\n"
        "3. 提出 2-3 个核心论点，每个论点配真实数据\n"
        "4. 论证要有力，语言要有说服力"
    ),
    tool_names=["web_search"],
)

DEBATE_CON = Agent(
    name="反方辩手", icon="🔴",
    system_prompt=(
        "你是反方辩手。针对给定辩题，你需要搜索资料并提出反对性论点。\n\n"
        "规则：\n"
        "1. 用 web_search 搜索真实案例和数据支持你的论点\n"
        "2. 每个论点要有事实依据，不能空谈\n"
        "3. 提出 2-3 个核心论点，每个论点配真实数据\n"
        "4. 要有力反驳正方观点，论证要犀利"
    ),
    tool_names=["web_search"],
)

DEBATE_JUDGE = Agent(
    name="裁判", icon="⚖️",
    system_prompt=(
        "你是中立裁判，综合正反双方的论据做出评判。\n\n"
        "规则：\n"
        "1. 用 calculator 或 run_python 做数据验证\n"
        "2. 逐点评判双方论据的强弱\n"
        "3. 给出结论，说明胜方理由\n"
        "4. 保持客观，不偏袒任何一方"
    ),
    tool_names=["calculator", "run_python"],
)

PIPELINE_SEARCHER = Agent(
    name="搜索 Agent", icon="🔍",
    system_prompt=(
        "你是信息搜索专家。\n\n"
        "规则：\n"
        "1. 用 web_search 搜索任务所需的信息\n"
        "2. 搜索多个角度，确保信息全面\n"
        "3. 返回结构化的搜索结果，标注来源\n"
        "4. 如果一次搜索不够，可以多次搜索"
    ),
    tool_names=["web_search"],
)

PIPELINE_ANALYZER = Agent(
    name="分析 Agent", icon="📊",
    system_prompt=(
        "你是数据分析专家。\n\n"
        "规则：\n"
        "1. 接收搜索 Agent 的结果\n"
        "2. 用 run_python 做数据处理和对比分析\n"
        "3. 用 calculator 做数值计算\n"
        "4. 输出结构化的分析结论"
    ),
    tool_names=["calculator", "run_python"],
)

PIPELINE_REPORTER = Agent(
    name="报告 Agent", icon="📝",
    system_prompt=(
        "你是报告撰写专家。\n\n"
        "规则：\n"
        "1. 接收分析 Agent 的结论\n"
        "2. 用 run_python 生成格式化报告\n"
        "3. 报告要有：标题、摘要、正文、结论、建议\n"
        "4. 语言专业但易读"
    ),
    tool_names=["run_python"],
)

# ── 执行引擎 ──

def _emit(emit, event_type, data):
    if emit:
        data["type"] = event_type
        data["timestamp"] = time.time()
        emit(data)


def run_supervisor(query: str, emit=None) -> dict:
    """Supervisor 模式：主管拆任务 → Researcher + Coder → 主管整合"""
    trace_id = str(uuid.uuid4())[:8]
    t_start = time.time()

    def agent_emit(event_type, data):
        data["trace_id"] = trace_id
        data["mode"] = "supervisor"
        _emit(emit, "agent_event", data)

    # Step 1: Supervisor 分析任务
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "Supervisor", "icon": "👔", "action": "分析任务并拆解"})
    supervisor_analysis = SUPERVISOR.chat(
        [{"role": "user", "content": f"任务：{query}\n\n请分析这个任务，拆成子任务，说明每个子任务交给哪个专家。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "Supervisor", "result": supervisor_analysis[:500]})

    # Step 2: Researcher 搜索信息
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "Researcher", "icon": "🔍", "action": "搜索信息"})
    research_result = RESEARCHER.chat(
        [{"role": "user", "content": f"请搜索以下任务所需的信息：\n{query}\n\n{supervisor_analysis}"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "Researcher", "result": research_result[:500]})

    # Step 3: Coder 分析计算
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "Coder", "icon": "💻", "action": "数据分析与计算"})
    coder_result = CODER.chat(
        [{"role": "user", "content": f"原始任务：{query}\n\n研究结果：\n{research_result}\n\n请基于以上数据做分析和计算。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "Coder", "result": coder_result[:500]})

    # Step 4: Supervisor 整合
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "Supervisor", "icon": "👔", "action": "整合最终答案"})
    final_answer = SUPERVISOR.chat(
        [{"role": "user", "content": (
            f"原始任务：{query}\n\n"
            f"研究专家的结果：\n{research_result}\n\n"
            f"数据分析的结果：\n{coder_result}\n\n"
            f"请整合以上信息，给出最终的完整答案。"
        )}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "Supervisor", "result": final_answer[:800]})

    total_ms = int((time.time() - t_start) * 1000)
    _emit(emit, "trace_complete", {"trace_id": trace_id, "mode": "supervisor", "total_ms": total_ms})

    return {"trace_id": trace_id, "answer": final_answer, "total_ms": total_ms}


def run_pipeline(query: str, emit=None) -> dict:
    """Pipeline 模式：搜索 → 分析 → 报告"""
    trace_id = str(uuid.uuid4())[:8]
    t_start = time.time()

    def agent_emit(event_type, data):
        data["trace_id"] = trace_id
        data["mode"] = "pipeline"
        _emit(emit, "agent_event", data)

    # Step 1: 搜索
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "搜索 Agent", "icon": "🔍", "action": "信息采集"})
    search_result = PIPELINE_SEARCHER.chat(
        [{"role": "user", "content": f"请搜索以下任务的信息：\n{query}"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "搜索 Agent", "result": search_result[:500]})

    # Step 2: 分析
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "分析 Agent", "icon": "📊", "action": "数据分析"})
    analysis_result = PIPELINE_ANALYZER.chat(
        [{"role": "user", "content": f"原始任务：{query}\n\n搜索结果：\n{search_result}\n\n请做数据分析。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "分析 Agent", "result": analysis_result[:500]})

    # Step 3: 生成报告
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "报告 Agent", "icon": "📝", "action": "生成报告"})
    report = PIPELINE_REPORTER.chat(
        [{"role": "user", "content": f"原始任务：{query}\n\n分析结果：\n{analysis_result}\n\n请生成最终报告。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "报告 Agent", "result": report[:800]})

    total_ms = int((time.time() - t_start) * 1000)
    _emit(emit, "trace_complete", {"trace_id": trace_id, "mode": "pipeline", "total_ms": total_ms})

    return {"trace_id": trace_id, "answer": report, "total_ms": total_ms}


def run_debate(query: str, emit=None) -> dict:
    """Debate 模式：正方 → 反方 → 二次交锋 → 裁判"""
    trace_id = str(uuid.uuid4())[:8]
    t_start = time.time()

    def agent_emit(event_type, data):
        data["trace_id"] = trace_id
        data["mode"] = "debate"
        _emit(emit, "agent_event", data)

    # Step 1: 正方立论
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "正方辩手", "icon": "🟢", "action": "立论陈词"})
    pro_round1 = DEBATE_PRO.chat(
        [{"role": "user", "content": f"辩题：{query}\n\n你是正方，请搜索资料并提出你的核心论点。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "正方辩手", "result": pro_round1[:500]})

    # Step 2: 反方反驳
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "反方辩手", "icon": "🔴", "action": "反驳 + 立论"})
    con_round1 = DEBATE_CON.chat(
        [{"role": "user", "content": f"辩题：{query}\n\n正方论点：\n{pro_round1}\n\n你是反方，请搜索资料反驳正方并提出你的论点。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "反方辩手", "result": con_round1[:500]})

    # Step 3: 正方回应
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "正方辩手", "icon": "🟢", "action": "回应反驳"})
    pro_round2 = DEBATE_PRO.chat(
        [{"role": "user", "content": f"辩题：{query}\n\n你的论点：\n{pro_round1}\n\n反方反驳：\n{con_round1}\n\n请回应反驳并强化你的论点。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "正方辩手", "result": pro_round2[:500]})

    # Step 4: 反方回应
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "反方辩手", "icon": "🔴", "action": "二次反驳"})
    con_round2 = DEBATE_CON.chat(
        [{"role": "user", "content": f"辩题：{query}\n\n你的论点：\n{con_round1}\n\n正方回应：\n{pro_round2}\n\n请做最终反驳。"}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "反方辩手", "result": con_round2[:500]})

    # Step 5: 裁判评判
    _emit(emit, "step_start", {"trace_id": trace_id, "agent": "裁判", "icon": "⚖️", "action": "综合评判"})
    verdict = DEBATE_JUDGE.chat(
        [{"role": "user", "content": (
            f"辩题：{query}\n\n"
            f"正方第一轮：\n{pro_round1}\n\n"
            f"反方第一轮：\n{con_round1}\n\n"
            f"正方第二轮：\n{pro_round2}\n\n"
            f"反方第二轮：\n{con_round2}\n\n"
            f"请综合评判，给出最终结论。"
        )}],
        emit=agent_emit
    )
    _emit(emit, "step_end", {"trace_id": trace_id, "agent": "裁判", "result": verdict[:800]})

    total_ms = int((time.time() - t_start) * 1000)
    _emit(emit, "trace_complete", {"trace_id": trace_id, "mode": "debate", "total_ms": total_ms})

    return {"trace_id": trace_id, "answer": verdict, "total_ms": total_ms}