"""ReWOO (Reasoning Without Observation Others) 执行引擎

实现 ReWOO 模式：先规划所有工具调用，再批量执行，最后综合结果。
与 ReAct 的区别在于：ReWOO 不在中间步骤观察结果来调整策略，
而是预先规划完整执行计划，一次性执行后综合输出。
"""

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

from agentkit.core.exceptions import TaskCancelledError, TaskTimeoutError
from agentkit.core.protocol import CancellationToken
from agentkit.core.react import ReActEngine, ReActEvent, ReActResult, ReActStep
from agentkit.llm.gateway import LLMGateway
from agentkit.llm.protocol import LLMResponse
from agentkit.tools.base import Tool
from agentkit.telemetry.tracing import get_tracer, start_span, _OTEL_AVAILABLE
from agentkit.telemetry.metrics import (
    agent_request_counter,
    agent_duration_histogram,
)

if TYPE_CHECKING:
    from agentkit.core.compressor import CompressionStrategy, ContextCompressor
    from agentkit.core.trace import TraceRecorder
    from agentkit.memory.retriever import MemoryRetriever

logger = logging.getLogger(__name__)


# ── Data Structures ───────────────────────────────────────


@dataclass
class ReWOOPlanStep:
    """ReWOO 计划中的单步"""

    step_id: int
    tool_name: str
    arguments: dict[str, Any]
    reasoning: str = ""


@dataclass
class ReWOOPlan:
    """ReWOO 执行计划"""

    steps: list[ReWOOPlanStep] = field(default_factory=list)
    reasoning: str = ""  # 整体规划推理


@dataclass
class ReWOOStep(ReActStep):
    """ReWOO 执行步骤，扩展 ReActStep 增加 plan_step_id"""

    plan_step_id: int | None = None


# ── Planning Prompt ───────────────────────────────────────

_PLANNING_SYSTEM_PROMPT = """\
You are a planning agent. Given a task and a set of available tools, \
create a step-by-step execution plan.

IMPORTANT: You must output a JSON object with the following structure:
{
  "reasoning": "Your overall reasoning about how to approach the task",
  "steps": [
    {
      "step_id": 1,
      "tool_name": "name_of_tool_to_call",
      "arguments": {"arg1": "value1", "arg2": "value2"},
      "reasoning": "Why this step is needed"
    },
    {
      "step_id": 2,
      "tool_name": "name_of_another_tool",
      "arguments": {"arg1": "value1"},
      "reasoning": "Why this step is needed"
    }
  ]
}

Rules:
- List ALL tool calls needed to complete the task in order
- Each step must use one of the available tools
- Arguments must match the tool's input schema
- If the task does not require any tools, return an empty steps list
- Output ONLY the JSON object, no other text
"""

_SYNTHESIS_SYSTEM_PROMPT = """\
You are a synthesis agent. Given the original task and the results of \
all planned tool executions, produce a final comprehensive answer.

Review all tool results below and synthesize them into a coherent response \
that fully addresses the original task.
"""


# ── ReWOO Engine ──────────────────────────────────────────


class ReWOOEngine:
    """ReWOO (Reasoning Without Observation Others) 执行引擎

    三阶段执行：
    1. Planning Phase: 一次性生成完整执行计划
    2. Execution Phase: 按计划顺序执行所有工具调用
    3. Synthesis Phase: 综合所有工具结果生成最终输出
    """

    FALLBACK_STRATEGIES = ["simplified_rewoo", "react", "direct"]

    def __init__(self, llm_gateway: LLMGateway, max_plan_steps: int = 10, default_timeout: float = 300.0):
        if max_plan_steps < 1:
            raise ValueError(f"max_plan_steps must be >= 1, got {max_plan_steps}")
        self._llm_gateway = llm_gateway
        self._max_plan_steps = max_plan_steps
        self._default_timeout = default_timeout
        # ReActEngine 作为 fallback
        self._react_engine = ReActEngine(
            llm_gateway=llm_gateway,
            max_steps=max_plan_steps,
            default_timeout=default_timeout,
        )

    async def execute(
        self,
        messages: list[dict[str, str]],
        tools: list[Tool] | None = None,
        model: str = "default",
        agent_name: str = "",
        task_type: str = "",
        system_prompt: str | None = None,
        trace_recorder: "TraceRecorder | None" = None,
        memory_retriever: "MemoryRetriever | None" = None,
        task_id: str | None = None,
        compressor: "CompressionStrategy | None" = None,
        retrieval_config: dict[str, Any] | None = None,
        cancellation_token: CancellationToken | None = None,
        timeout_seconds: float | None = None,
    ) -> ReActResult:
        """执行 ReWOO 三阶段流程

        1. Planning: 调用 LLM 生成完整执行计划
        2. Execution: 按计划顺序执行所有工具调用
        3. Synthesis: 调用 LLM 综合所有结果生成最终输出

        如果 Planning 阶段失败（LLM 未返回有效 JSON），则回退到 ReActEngine。

        Args:
            cancellation_token: 协作式取消令牌
            timeout_seconds: 超时秒数，0 表示无超时，None 使用 default_timeout
        """
        effective_timeout = timeout_seconds if timeout_seconds is not None else self._default_timeout

        try:
            if effective_timeout > 0:
                result = await asyncio.wait_for(
                    self._execute_rewoo(
                        messages=messages,
                        tools=tools,
                        model=model,
                        agent_name=agent_name,
                        task_type=task_type,
                        system_prompt=system_prompt,
                        trace_recorder=trace_recorder,
                        memory_retriever=memory_retriever,
                        task_id=task_id,
                        compressor=compressor,
                        retrieval_config=retrieval_config,
                        cancellation_token=cancellation_token,
                    ),
                    timeout=effective_timeout,
                )
            else:
                result = await self._execute_rewoo(
                    messages=messages,
                    tools=tools,
                    model=model,
                    agent_name=agent_name,
                    task_type=task_type,
                    system_prompt=system_prompt,
                    trace_recorder=trace_recorder,
                    memory_retriever=memory_retriever,
                    task_id=task_id,
                    compressor=compressor,
                    retrieval_config=retrieval_config,
                    cancellation_token=cancellation_token,
                )
        except asyncio.TimeoutError:
            raise TaskTimeoutError(
                task_id=task_id or "",
                timeout_seconds=int(effective_timeout),
            )
        except TaskCancelledError:
            raise

        return result

    async def _execute_rewoo(
        self,
        messages: list[dict[str, str]],
        tools: list[Tool] | None = None,
        model: str = "default",
        agent_name: str = "",
        task_type: str = "",
        system_prompt: str | None = None,
        trace_recorder: "TraceRecorder | None" = None,
        memory_retriever: "MemoryRetriever | None" = None,
        task_id: str | None = None,
        compressor: "CompressionStrategy | None" = None,
        retrieval_config: dict[str, Any] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> ReActResult:
        tools = tools or []
        tool_schemas = self._build_tool_schemas(tools) if tools else None

        # Telemetry: record agent request
        agent_request_counter().add(1, {"agent.name": agent_name, "agent.type": task_type or "rewoo"})

        # Start telemetry span
        _span_cm = None
        _span = None
        _exec_start = time.monotonic()

        if _OTEL_AVAILABLE:
            _span_cm = start_span(
                "agent.execute.rewoo",
                attributes={"agent.name": agent_name, "agent.type": task_type or "rewoo"},
            )
            _span = _span_cm.__enter__()

        # Initialize before try so finally can access them
        trajectory: list[ReActStep] = []
        total_tokens = 0
        trace_outcome = "error"

        try:
            # 启动轨迹记录
            if trace_recorder is not None:
                trace_recorder.start_trace(
                    task_id="",
                    agent_name=agent_name,
                    skill_name=task_type or None,
                )

            # Memory retrieval: 执行前检索相关上下文注入 system_prompt
            effective_system_prompt = system_prompt
            if memory_retriever:
                try:
                    query = str(messages[-1].get("content", "")) if messages else ""
                    top_k = (retrieval_config or {}).get("top_k", 5)
                    token_budget = (retrieval_config or {}).get("token_budget", 2000)
                    memory_context = await memory_retriever.get_context_string(
                        query=query,
                        top_k=top_k,
                        token_budget=token_budget,
                    )
                    if memory_context:
                        if effective_system_prompt:
                            effective_system_prompt += f"\n\n## 参考信息\n{memory_context}"
                        else:
                            effective_system_prompt = f"## 参考信息\n{memory_context}"
                except Exception as e:
                    logger.warning(f"Memory retrieval failed, continuing without context: {e}")

            # ── Phase 1: Planning ──
            plan, planning_tokens = await self._plan_phase(
                messages=messages,
                tools=tools,
                tool_schemas=tool_schemas,
                model=model,
                agent_name=agent_name,
                task_type=task_type,
                system_prompt=effective_system_prompt,
                compressor=compressor,
                cancellation_token=cancellation_token,
            )
            total_tokens += planning_tokens

            fallback_strategy: str | None = None

            # 记录规划步骤
            if trace_recorder is not None:
                trace_recorder.record_step(
                    step=0,
                    action="planning",
                    duration_ms=0,
                    tokens_used=planning_tokens,
                )

            # 如果规划失败，尝试渐进式回退
            if plan is None:
                # 尝试简化规划（max_steps=3）
                logger.warning("ReWOO planning failed, trying simplified planning with max_steps=3")
                try:
                    plan, simplified_tokens = await self._plan_phase(
                        messages=messages,
                        tools=tools,
                        tool_schemas=tool_schemas,
                        model=model,
                        agent_name=agent_name,
                        task_type=task_type,
                        system_prompt=effective_system_prompt,
                        compressor=compressor,
                        cancellation_token=cancellation_token,
                        max_steps=3,
                    )
                    total_tokens += simplified_tokens
                    if plan is not None and plan.steps:
                        fallback_strategy = "simplified_rewoo"
                        logger.info("Simplified ReWOO planning succeeded")
                except Exception as e2:
                    logger.warning(f"Simplified ReWOO planning also failed: {e2}")

            if plan is None:
                # 回退到 ReAct
                fallback_strategy = "react"
                logger.warning("ReWOO planning failed, falling back to ReActEngine")
                if trace_recorder is not None:
                    trace_recorder.end_trace(outcome="fallback")
                try:
                    react_result = await self._react_engine.execute(
                        messages=messages,
                        tools=tools,
                        model=model,
                        agent_name=agent_name,
                        task_type=task_type,
                        system_prompt=system_prompt,
                        trace_recorder=trace_recorder,
                        memory_retriever=memory_retriever,
                        task_id=task_id,
                        compressor=compressor,
                        retrieval_config=retrieval_config,
                        cancellation_token=cancellation_token,
                        timeout_seconds=0,  # timeout already handled by outer wrapper
                    )
                    react_result.fallback_strategy = fallback_strategy
                    return react_result
                except Exception as react_err:
                    # ReAct 也失败，回退到 Direct（简单 LLM 调用）
                    fallback_strategy = "direct"
                    logger.warning(f"ReAct fallback also failed: {react_err}, falling back to direct LLM call")
                    try:
                        direct_messages: list[dict[str, Any]] = []
                        if effective_system_prompt:
                            direct_messages.append({"role": "system", "content": effective_system_prompt})
                        direct_messages.extend(messages)

                        if compressor:
                            try:
                                direct_messages = await compressor.compress(direct_messages)
                            except Exception as e:
                                logger.warning(f"Context compression failed in direct fallback: {e}")

                        direct_response = await self._llm_gateway.chat(
                            messages=direct_messages,
                            model=model,
                            agent_name=agent_name,
                            task_type=task_type,
                        )
                        total_tokens += direct_response.usage.total_tokens

                        direct_step = ReWOOStep(
                            step=1,
                            action="final_answer",
                            content=direct_response.content,
                            tokens=direct_response.usage.total_tokens,
                            plan_step_id=None,
                        )
                        trajectory.append(direct_step)

                        if trace_recorder is not None:
                            trace_recorder.record_step(
                                step=1,
                                action="final_answer",
                                output_data={"content": direct_response.content},
                                tokens_used=direct_response.usage.total_tokens,
                            )

                        trace_outcome = "success"
                        if trace_recorder is not None:
                            trace_recorder.end_trace(outcome=trace_outcome)

                        return ReActResult(
                            output=direct_response.content or "",
                            trajectory=trajectory,
                            total_steps=len(trajectory),
                            total_tokens=total_tokens,
                            fallback_strategy=fallback_strategy,
                        )
                    except Exception as direct_err:
                        logger.error(f"Direct LLM fallback also failed: {direct_err}")
                        raise

            # 如果计划为空（无需工具），直接让 LLM 回答
            if not plan.steps:
                llm_messages: list[dict[str, Any]] = []
                if effective_system_prompt:
                    llm_messages.append({"role": "system", "content": effective_system_prompt})
                llm_messages.extend(messages)

                if compressor:
                    try:
                        llm_messages = await compressor.compress(llm_messages)
                    except Exception as e:
                        logger.warning(f"Context compression failed: {e}")

                response = await self._llm_gateway.chat(
                    messages=llm_messages,
                    model=model,
                    agent_name=agent_name,
                    task_type=task_type,
                )
                total_tokens += response.usage.total_tokens

                step = ReWOOStep(
                    step=1,
                    action="final_answer",
                    content=response.content,
                    tokens=response.usage.total_tokens,
                    plan_step_id=None,
                )
                trajectory.append(step)

                if trace_recorder is not None:
                    trace_recorder.record_step(
                        step=1,
                        action="final_answer",
                        output_data={"content": response.content},
                        tokens_used=response.usage.total_tokens,
                    )

                trace_outcome = "success"
                if trace_recorder is not None:
                    trace_recorder.end_trace(outcome=trace_outcome)

                return ReActResult(
                    output=response.content or "",
                    trajectory=trajectory,
                    total_steps=len(trajectory),
                    total_tokens=total_tokens,
                    fallback_strategy=fallback_strategy,
                )

            # ── Phase 2: Execution ──
            tool_results: list[dict[str, Any]] = []
            for plan_step in plan.steps:
                # 协作式取消检查
                if cancellation_token is not None:
                    cancellation_token.check()

                tool_start = time.monotonic()
                tool_result = await self._execute_tool(plan_step.tool_name, plan_step.arguments, tools)
                tool_duration_ms = int((time.monotonic() - tool_start) * 1000)

                rewoo_step = ReWOOStep(
                    step=plan_step.step_id,
                    action="tool_call",
                    tool_name=plan_step.tool_name,
                    arguments=plan_step.arguments,
                    result=tool_result,
                    tokens=0,  # tool execution tokens tracked separately
                    plan_step_id=plan_step.step_id,
                )
                trajectory.append(rewoo_step)

                tool_results.append({
                    "step_id": plan_step.step_id,
                    "tool_name": plan_step.tool_name,
                    "arguments": plan_step.arguments,
                    "result": tool_result,
                    "reasoning": plan_step.reasoning,
                })

                # 记录工具调用步骤
                if trace_recorder is not None:
                    tool_error = None
                    if isinstance(tool_result, dict) and "error" in tool_result:
                        tool_error = tool_result["error"]
                    trace_recorder.record_step(
                        step=plan_step.step_id,
                        action="tool_call",
                        tool_name=plan_step.tool_name,
                        input_data=plan_step.arguments,
                        output_data=tool_result,
                        duration_ms=tool_duration_ms,
                        tokens_used=0,
                        error=tool_error,
                    )

            # ── Phase 3: Synthesis ──
            output, synthesis_tokens = await self._synthesis_phase(
                messages=messages,
                tool_results=tool_results,
                model=model,
                agent_name=agent_name,
                task_type=task_type,
                system_prompt=effective_system_prompt,
                compressor=compressor,
                cancellation_token=cancellation_token,
            )
            total_tokens += synthesis_tokens

            # 记录综合步骤
            synthesis_step = ReWOOStep(
                step=len(plan.steps) + 1,
                action="final_answer",
                content=output,
                tokens=synthesis_tokens,
                plan_step_id=None,
            )
            trajectory.append(synthesis_step)

            if trace_recorder is not None:
                trace_recorder.record_step(
                    step=len(plan.steps) + 1,
                    action="final_answer",
                    output_data={"content": output},
                    tokens_used=synthesis_tokens,
                )

            trace_outcome = "success"

            # 结束轨迹记录
            if trace_recorder is not None:
                trace_recorder.end_trace(outcome=trace_outcome)

            # Memory storage: 执行后写入轨迹摘要到 EpisodicMemory
            if memory_retriever and hasattr(memory_retriever, "store_episode"):
                try:
                    summary = output[:500] if output else ""
                    await memory_retriever.store_episode(
                        key=f"task:{task_id or 'unknown'}",
                        value={"output_summary": summary, "agent_name": agent_name},
                        metadata={"task_type": task_type, "outcome": trace_outcome},
                    )
                except Exception as e:
                    logger.warning(f"Failed to store task result in episodic memory: {e}")

            return ReActResult(
                output=output,
                trajectory=trajectory,
                total_steps=len(trajectory),
                total_tokens=total_tokens,
                fallback_strategy=fallback_strategy,
            )
        finally:
            # Telemetry: end span and record duration
            _duration_ms = int((time.monotonic() - _exec_start) * 1000)
            if _span is not None:
                _span.set_attribute("agent.total_steps", len(trajectory))
                _span.set_attribute("agent.total_tokens", total_tokens)
                _span.set_attribute("agent.outcome", trace_outcome)
                _span.set_attribute("agent.duration_ms", _duration_ms)
            if _span_cm is not None:
                _span_cm.__exit__(None, None, None)
            agent_duration_histogram().record(_duration_ms, {"agent.name": agent_name})

    async def execute_stream(
        self,
        messages: list[dict[str, str]],
        tools: list[Tool] | None = None,
        model: str = "default",
        agent_name: str = "",
        task_type: str = "",
        system_prompt: str | None = None,
        trace_recorder: "TraceRecorder | None" = None,
        memory_retriever: "MemoryRetriever | None" = None,
        task_id: str | None = None,
        compressor: "CompressionStrategy | None" = None,
        retrieval_config: dict[str, Any] | None = None,
        cancellation_token: CancellationToken | None = None,
        timeout_seconds: float | None = None,
    ):
        """Execute ReWOO flow, yielding ReActEvent objects.

        Events:
        - "planning": planning phase started
        - "plan_generated": plan generated with step details
        - "tool_call": a tool is being called
        - "tool_result": tool execution result
        - "synthesis": synthesis phase started
        - "final_answer": final synthesized answer
        """
        tools = tools or []
        tool_schemas = self._build_tool_schemas(tools) if tools else None

        # 启动轨迹记录
        if trace_recorder is not None:
            trace_recorder.start_trace(
                task_id="",
                agent_name=agent_name,
                skill_name=task_type or None,
            )

        # Memory retrieval
        effective_system_prompt = system_prompt
        if memory_retriever:
            try:
                query = str(messages[-1].get("content", "")) if messages else ""
                top_k = (retrieval_config or {}).get("top_k", 5)
                token_budget = (retrieval_config or {}).get("token_budget", 2000)
                memory_context = await memory_retriever.get_context_string(
                    query=query,
                    top_k=top_k,
                    token_budget=token_budget,
                )
                if memory_context:
                    if effective_system_prompt:
                        effective_system_prompt += f"\n\n## 参考信息\n{memory_context}"
                    else:
                        effective_system_prompt = f"## 参考信息\n{memory_context}"
            except Exception as e:
                logger.warning(f"Memory retrieval failed, continuing without context: {e}")

        trajectory: list[ReActStep] = []
        total_tokens = 0
        output = ""
        trace_outcome = "success"

        try:
            # ── Phase 1: Planning ──
            yield ReActEvent(
                event_type="planning",
                step=0,
                data={"message": "Generating execution plan..."},
            )

            plan, planning_tokens = await self._plan_phase(
                messages=messages,
                tools=tools,
                tool_schemas=tool_schemas,
                model=model,
                agent_name=agent_name,
                task_type=task_type,
                system_prompt=effective_system_prompt,
                compressor=compressor,
                cancellation_token=cancellation_token,
            )
            total_tokens += planning_tokens

            if plan is None:
                # Try simplified planning
                logger.warning("ReWOO planning failed in stream mode, trying simplified planning with max_steps=3")
                try:
                    plan, simplified_tokens = await self._plan_phase(
                        messages=messages,
                        tools=tools,
                        tool_schemas=tool_schemas,
                        model=model,
                        agent_name=agent_name,
                        task_type=task_type,
                        system_prompt=effective_system_prompt,
                        compressor=compressor,
                        cancellation_token=cancellation_token,
                        max_steps=3,
                    )
                    total_tokens += simplified_tokens
                except Exception as e2:
                    logger.warning(f"Simplified ReWOO planning also failed in stream mode: {e2}")

            if plan is None:
                # Planning failed, fall back to ReAct streaming
                logger.warning("ReWOO planning failed in stream mode, falling back to ReActEngine")
                try:
                    async for event in self._react_engine.execute_stream(
                        messages=messages,
                        tools=tools,
                        model=model,
                        agent_name=agent_name,
                        task_type=task_type,
                        system_prompt=system_prompt,
                        trace_recorder=trace_recorder,
                        memory_retriever=memory_retriever,
                        task_id=task_id,
                        compressor=compressor,
                        retrieval_config=retrieval_config,
                        cancellation_token=cancellation_token,
                        timeout_seconds=0,
                    ):
                        yield event
                    return
                except Exception as react_err:
                    # ReAct also failed, fall back to direct LLM call
                    logger.warning(f"ReAct fallback also failed in stream mode: {react_err}, falling back to direct LLM call")
                    try:
                        direct_messages: list[dict[str, Any]] = []
                        if effective_system_prompt:
                            direct_messages.append({"role": "system", "content": effective_system_prompt})
                        direct_messages.extend(messages)

                        if compressor:
                            try:
                                direct_messages = await compressor.compress(direct_messages)
                            except Exception as e:
                                logger.warning(f"Context compression failed in direct fallback: {e}")

                        direct_response = await self._llm_gateway.chat(
                            messages=direct_messages,
                            model=model,
                            agent_name=agent_name,
                            task_type=task_type,
                        )
                        total_tokens += direct_response.usage.total_tokens
                        output = direct_response.content or ""

                        yield ReActEvent(
                            event_type="final_answer",
                            step=1,
                            data={
                                "output": output,
                                "total_steps": 1,
                                "total_tokens": total_tokens,
                            },
                        )
                        return
                    except Exception as direct_err:
                        logger.error(f"Direct LLM fallback also failed in stream mode: {direct_err}")
                        raise

            yield ReActEvent(
                event_type="plan_generated",
                step=0,
                data={
                    "reasoning": plan.reasoning,
                    "steps": [
                        {
                            "step_id": s.step_id,
                            "tool_name": s.tool_name,
                            "arguments": s.arguments,
                            "reasoning": s.reasoning,
                        }
                        for s in plan.steps
                    ],
                },
            )

            # Empty plan: direct answer
            if not plan.steps:
                llm_messages: list[dict[str, Any]] = []
                if effective_system_prompt:
                    llm_messages.append({"role": "system", "content": effective_system_prompt})
                llm_messages.extend(messages)

                if compressor:
                    try:
                        llm_messages = await compressor.compress(llm_messages)
                    except Exception as e:
                        logger.warning(f"Context compression failed: {e}")

                response = await self._llm_gateway.chat(
                    messages=llm_messages,
                    model=model,
                    agent_name=agent_name,
                    task_type=task_type,
                )
                total_tokens += response.usage.total_tokens
                output = response.content or ""

                trajectory.append(ReWOOStep(
                    step=1,
                    action="final_answer",
                    content=output,
                    tokens=response.usage.total_tokens,
                ))

                yield ReActEvent(
                    event_type="final_answer",
                    step=1,
                    data={
                        "output": output,
                        "total_steps": len(trajectory),
                        "total_tokens": total_tokens,
                    },
                )
                return

            # ── Phase 2: Execution ──
            tool_results: list[dict[str, Any]] = []
            for plan_step in plan.steps:
                if cancellation_token is not None:
                    cancellation_token.check()

                yield ReActEvent(
                    event_type="tool_call",
                    step=plan_step.step_id,
                    data={"tool_name": plan_step.tool_name, "arguments": plan_step.arguments},
                )

                tool_start = time.monotonic()
                tool_result = await self._execute_tool(plan_step.tool_name, plan_step.arguments, tools)
                tool_duration_ms = int((time.monotonic() - tool_start) * 1000)

                rewoo_step = ReWOOStep(
                    step=plan_step.step_id,
                    action="tool_call",
                    tool_name=plan_step.tool_name,
                    arguments=plan_step.arguments,
                    result=tool_result,
                    tokens=0,
                    plan_step_id=plan_step.step_id,
                )
                trajectory.append(rewoo_step)

                tool_results.append({
                    "step_id": plan_step.step_id,
                    "tool_name": plan_step.tool_name,
                    "arguments": plan_step.arguments,
                    "result": tool_result,
                    "reasoning": plan_step.reasoning,
                })

                # 记录工具调用步骤
                if trace_recorder is not None:
                    tool_error = None
                    if isinstance(tool_result, dict) and "error" in tool_result:
                        tool_error = tool_result["error"]
                    trace_recorder.record_step(
                        step=plan_step.step_id,
                        action="tool_call",
                        tool_name=plan_step.tool_name,
                        input_data=plan_step.arguments,
                        output_data=tool_result,
                        duration_ms=tool_duration_ms,
                        tokens_used=0,
                        error=tool_error,
                    )

                yield ReActEvent(
                    event_type="tool_result",
                    step=plan_step.step_id,
                    data={"tool_name": plan_step.tool_name, "result": tool_result},
                )

            # ── Phase 3: Synthesis ──
            yield ReActEvent(
                event_type="synthesis",
                step=len(plan.steps) + 1,
                data={"message": "Synthesizing results..."},
            )

            output, synthesis_tokens = await self._synthesis_phase(
                messages=messages,
                tool_results=tool_results,
                model=model,
                agent_name=agent_name,
                task_type=task_type,
                system_prompt=effective_system_prompt,
                compressor=compressor,
                cancellation_token=cancellation_token,
            )
            total_tokens += synthesis_tokens

            trajectory.append(ReWOOStep(
                step=len(plan.steps) + 1,
                action="final_answer",
                content=output,
                tokens=synthesis_tokens,
            ))

            yield ReActEvent(
                event_type="final_answer",
                step=len(plan.steps) + 1,
                data={
                    "output": output,
                    "total_steps": len(trajectory),
                    "total_tokens": total_tokens,
                },
            )
        finally:
            # 结束轨迹记录
            if trace_recorder is not None:
                trace_recorder.end_trace(outcome=trace_outcome)

            # Memory storage
            if memory_retriever and hasattr(memory_retriever, "store_episode"):
                try:
                    summary = output[:500] if output else ""
                    await memory_retriever.store_episode(
                        key=f"task:{task_id or 'unknown'}",
                        value={"output_summary": summary, "agent_name": agent_name},
                        metadata={"task_type": task_type, "outcome": trace_outcome},
                    )
                except Exception as e:
                    logger.warning(f"Failed to store task result in episodic memory: {e}")

    # ── Phase Implementations ─────────────────────────────

    async def _plan_phase(
        self,
        messages: list[dict[str, str]],
        tools: list[Tool],
        tool_schemas: list[dict] | None,
        model: str,
        agent_name: str,
        task_type: str,
        system_prompt: str | None,
        compressor: "CompressionStrategy | None",
        cancellation_token: CancellationToken | None,
        max_steps: int | None = None,
    ) -> tuple[ReWOOPlan | None, int]:
        """Planning Phase: 调用 LLM 生成完整执行计划

        Args:
            max_steps: 限制计划最大步数，None 则使用 self._max_plan_steps

        Returns:
            (plan, tokens_used) - plan 为 None 表示规划失败
        """
        if cancellation_token is not None:
            cancellation_token.check()

        # 构建工具描述
        tool_descriptions = self._build_tool_descriptions(tools)

        # 构建规划消息
        planning_messages: list[dict[str, Any]] = [
            {"role": "system", "content": _PLANNING_SYSTEM_PROMPT},
        ]

        # 添加上下文信息
        context_parts = []
        if system_prompt:
            context_parts.append(f"Context: {system_prompt}")
        if tool_descriptions:
            context_parts.append(f"Available tools:\n{tool_descriptions}")

        user_content = "\n\n".join(context_parts) if context_parts else ""
        # 添加原始用户消息
        for msg in messages:
            if msg.get("role") == "user":
                user_content += f"\n\nTask: {msg.get('content', '')}"

        planning_messages.append({"role": "user", "content": user_content})

        # 压缩
        if compressor:
            try:
                planning_messages = await compressor.compress(planning_messages)
            except Exception as e:
                logger.warning(f"Context compression failed during planning: {e}")

        try:
            response = await self._llm_gateway.chat(
                messages=planning_messages,
                model=model,
                agent_name=agent_name,
                task_type=task_type,
                tools=tool_schemas,
            )
        except Exception as e:
            logger.warning(f"LLM call failed during planning: {e}")
            return None, 0

        tokens_used = response.usage.total_tokens

        # 解析计划
        plan = self._parse_plan(response.content or "")
        if plan is None:
            return None, tokens_used

        # 限制计划步数
        effective_max_steps = max_steps if max_steps is not None else self._max_plan_steps
        if len(plan.steps) > effective_max_steps:
            plan.steps = plan.steps[:effective_max_steps]

        return plan, tokens_used

    async def _synthesis_phase(
        self,
        messages: list[dict[str, str]],
        tool_results: list[dict[str, Any]],
        model: str,
        agent_name: str,
        task_type: str,
        system_prompt: str | None,
        compressor: "CompressionStrategy | None",
        cancellation_token: CancellationToken | None,
    ) -> tuple[str, int]:
        """Synthesis Phase: 综合所有工具结果生成最终输出

        Returns:
            (output, tokens_used)
        """
        if cancellation_token is not None:
            cancellation_token.check()

        # 构建综合消息
        synthesis_messages: list[dict[str, Any]] = [
            {"role": "system", "content": _SYNTHESIS_SYSTEM_PROMPT},
        ]

        # 构建工具结果摘要
        results_text = "Tool execution results:\n\n"
        for tr in tool_results:
            results_text += f"Step {tr['step_id']}: {tr['tool_name']}"
            if tr.get("reasoning"):
                results_text += f" (Reason: {tr['reasoning']})"
            results_text += "\n"
            results_text += f"  Arguments: {json.dumps(tr['arguments'], ensure_ascii=False)}\n"
            results_text += f"  Result: {json.dumps(tr['result'], ensure_ascii=False, default=str)}\n\n"

        # 添加原始用户消息
        user_content = results_text
        for msg in messages:
            if msg.get("role") == "user":
                user_content = f"Original task: {msg.get('content', '')}\n\n{user_content}"

        if system_prompt:
            user_content = f"Context: {system_prompt}\n\n{user_content}"

        synthesis_messages.append({"role": "user", "content": user_content})

        # 压缩
        if compressor:
            try:
                synthesis_messages = await compressor.compress(synthesis_messages)
            except Exception as e:
                logger.warning(f"Context compression failed during synthesis: {e}")

        response = await self._llm_gateway.chat(
            messages=synthesis_messages,
            model=model,
            agent_name=agent_name,
            task_type=task_type,
        )

        return response.content or "", response.usage.total_tokens

    # ── Helper Methods ────────────────────────────────────

    def _build_tool_schemas(self, tools: list[Tool]) -> list[dict]:
        """将 Tool 对象转换为 OpenAI Function Calling schema 格式"""
        schemas = []
        for tool in tools:
            schema = {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.input_schema or {"type": "object", "properties": {}},
                },
            }
            schemas.append(schema)
        return schemas

    def _build_tool_descriptions(self, tools: list[Tool]) -> str:
        """构建工具描述文本，用于规划 prompt"""
        descriptions = []
        for tool in tools:
            desc = f"- {tool.name}: {tool.description}"
            if tool.input_schema:
                props = tool.input_schema.get("properties", {})
                if props:
                    params = ", ".join(
                        f"{k} ({v.get('type', 'any')}: {v.get('description', '')})"
                        for k, v in props.items()
                    )
                    desc += f"\n  Parameters: {params}"
            descriptions.append(desc)
        return "\n".join(descriptions)

    def _parse_plan(self, content: str) -> ReWOOPlan | None:
        """从 LLM 响应中解析执行计划

        尝试从响应内容中提取 JSON 格式的计划。
        支持纯 JSON 和 markdown 代码块中的 JSON。
        """
        # 尝试提取 JSON 代码块
        json_str = content.strip()

        # 尝试从 markdown 代码块中提取
        if "```" in json_str:
            import re
            code_block_match = re.search(r"```(?:json)?\s*\n(.*?)\n\s*```", json_str, re.DOTALL)
            if code_block_match:
                json_str = code_block_match.group(1).strip()

        # 尝试提取 JSON 对象（处理 LLM 可能在 JSON 前后添加文本的情况）
        brace_start = json_str.find("{")
        brace_end = json_str.rfind("}")
        if brace_start != -1 and brace_end != -1 and brace_end > brace_start:
            json_str = json_str[brace_start:brace_end + 1]

        try:
            data = json.loads(json_str)
        except (json.JSONDecodeError, TypeError):
            logger.warning(f"Failed to parse plan from LLM response: {content[:200]}")
            return None

        if not isinstance(data, dict) or "steps" not in data:
            logger.warning(f"Plan JSON missing 'steps' key: {content[:200]}")
            return None

        steps = []
        for i, step_data in enumerate(data["steps"]):
            if not isinstance(step_data, dict):
                continue
            tool_name = step_data.get("tool_name", "")
            if not tool_name:
                continue
            steps.append(ReWOOPlanStep(
                step_id=step_data.get("step_id", i + 1),
                tool_name=tool_name,
                arguments=step_data.get("arguments", {}),
                reasoning=step_data.get("reasoning", ""),
            ))

        return ReWOOPlan(
            steps=steps,
            reasoning=data.get("reasoning", ""),
        )

    def _find_tool(self, name: str, tools: list[Tool]) -> Tool | None:
        """根据名称从可用工具中查找工具"""
        for tool in tools:
            if tool.name == name:
                return tool
        return None

    async def _execute_tool(
        self, tool_name: str, arguments: dict[str, Any], tools: list[Tool]
    ) -> dict:
        """执行工具调用，处理成功和失败情况"""
        tool = self._find_tool(tool_name, tools)
        if tool is None:
            error_msg = f"Tool '{tool_name}' not found"
            logger.warning(error_msg)
            return {"error": error_msg}

        try:
            result = await tool.safe_execute(**arguments)
            return result
        except Exception as e:
            error_msg = f"Tool '{tool_name}' execution failed: {e}"
            logger.warning(error_msg)
            return {"error": error_msg}
