import logging
import re
import time
import traceback
from copy import deepcopy
from types import GeneratorType
from typing import Callable, Iterable, Iterator, List, Optional, Type, TypeVar, Union

from lxml import etree
from pydantic import BaseModel, Field

from duowen_agent.llm import MessagesSet
from duowen_agent.tools.base import BaseTool, BaseToolResult
from duowen_agent.tools.manager import ToolManager
from duowen_agent.tools.mcp_client import McpClient
from duowen_agent.utils.core_utils import stream_to_string
from duowen_agent.utils.string_template import StringTemplate
from duowen_agent.utils.xml_tool_parser import parse_xml_tool_calls
from .base import BaseAgent
from .state import Resources
from ..error import ObserverException
from ..llm.chat_model import BaseAIChat
from ..utils.core_utils import generate_unique_id, remove_think

REACT_SYSTEM_PROMPT_TEMPLATE = StringTemplate(
    template_format="jinja2",
    template="""作为勤勉的任务代理，你的目标是尽可能高效地完成提供的任务或问题。

## Tools
你可以使用以下工具，工具信息通过以下结构描述提供：
{{tool_descriptions}}

## Output Format
请使用以下XML格式回答问题。仅返回XML，不要有解释。否则将会受到惩罚。
输出必须符合以下XML格式规范。仅返回XML，不要解释。

```xml
<response>
  <analysis>当前操作思路及原因分析</analysis>
  <function_calls>
    <invoke name="工具名称">
      <parameter name="参数名1">参数值1</parameter>
      <parameter name="参数名2">["参数值1","参数值2"]</parameter>
    </invoke>
  </function_calls>
</response>
```

如果使用此格式，用户将按以下格式回应：

<function_calls_result>
  <name>工具名称</name>
  <result>调用结果</result>
</function_calls_result>


你应持续重复上述格式，同时每次Observation进行反思，直到获得足够信息无需继续使用工具即可回答问题。此时必须改用以下两种格式之一：

- 若能回答问题：

```xml
<response>
  <analysis>最终判断依据及结论</analysis>
  <function_calls>
    <invoke name="final_answer">
      <parameter name="content">你的答案</parameter>
    </invoke>
  </function_calls>
</response>
```

- 若当前上下文无法回答问题：

```xml
<response>
  <analysis>无法解决的原因分析</analysis>
  <function_calls>
    <invoke name="final_answer">
      <parameter name="content">抱歉，我无法回答您的问题，因为（总结所有步骤并说明原因）</parameter>
    </invoke>
  </function_calls>
</response>
```

## Attention
- 仅返回XML格式，不要有解释
- 每一步只能选择一个工具
- 最终答案语言需与用户提问语言一致
- 动作参数的格式（XML或字符串）取决于工具定义
""",  # noqa: E501
)


class ActionResponseArgs(BaseModel):
    name: str
    args: Union[dict, str]


class ActionResponse(BaseModel):
    analysis: str
    action: ActionResponseArgs


class ReactActionParameters(BaseModel):
    analysis: Optional[str] = None
    action_name: str
    action_parameters: Optional[dict] = None


class ReactAction(BaseModel):
    action: ReactActionParameters
    cell_id: str
    call_id: str
    agent_name: Optional[str] = None


T = TypeVar("T")


class ReactObservationResult(BaseModel):
    result: str
    view: Optional[T] = None


class ReactObservation(BaseModel):
    observation: ReactObservationResult
    action: ReactActionParameters
    exec_status: bool
    cell_id: str
    call_id: str
    agent_name: Optional[str] = None


class ReactResult(BaseModel):
    analysis: str
    result: str | int | float
    cell_id: str
    agent_name: Optional[str] = None


class AgentToolParameters(BaseModel):
    instruction: str = Field(..., description="给子智能体的指令")


class AgentToolWrapper(BaseTool):
    name: str
    description: str
    parameters: Type[BaseModel] = AgentToolParameters

    def __init__(self, agent: "ReactAgent", name: str, description: str, **kwargs):
        self.agent = agent
        self.name = name
        self.description = description
        super().__init__(**kwargs)

    def _run(self, instruction: str):
        """
        运行子智能体，并打印中间过程
        """
        # 记录最后的结果
        final_result = ""
        try:
            # 消费生成器
            for item in self.agent.run(instruction):
                yield item
                if isinstance(item, ReactResult):
                    final_result = item.result

            # 如果循环结束 final_result 仍为空，说明没有找到最终答案
            if not final_result:
                error_msg = (
                    f"SubAgent {self.name} executed but did not provide a final answer."
                )
                logging.warning(error_msg)
                # return error_msg # Generator don't return values like this

        except Exception as e:
            error_msg = f"SubAgent execution failed: {str(e)}"
            logging.error(error_msg)
            # return error_msg


class ReactAgent(BaseAgent):
    def __init__(
        self,
        llm: BaseAIChat,
        tools: List[BaseTool] = None,
        subs_agent: List["ReactAgent"] = None,
        description: str = None,
        mcp_client: Optional[McpClient] = None,
        prefix_prompt: str = None,
        filter_function_list: List[str] = None,
        max_iterations: Optional[int] = 15,
        is_interrupt: Callable[[], bool] = None,  # 外部中断函数
        resources: Resources = None,
        agent_name: Optional[str] = None,
    ):
        super().__init__()
        # 处理 subs_agent，将其转换为工具
        if subs_agent:
            if tools is None:
                tools = []
            for sa in subs_agent:
                # 使用传入的 description 或者 agent 自带的 description (如果 ReactAgent 有这个属性的话，这里假设没有则用 name)
                # 注意：ReactAgent 本身没有 description 属性，所以这里优先使用参数中的，或者 fallback 到 name
                sa_desc = (
                    getattr(sa, "description", None)
                    or f"一个名为 {sa.agent_name} 的智能助手"
                )
                tools.append(
                    AgentToolWrapper(
                        agent=sa,
                        name=sa.agent_name,
                        description=sa_desc,
                    )
                )

        self.tool_manager = ToolManager(tools, filter_function_list) if tools else None
        self.mcp_client = mcp_client
        self.llm = llm
        self.description = description
        self.system_prompt_template: StringTemplate = REACT_SYSTEM_PROMPT_TEMPLATE
        self.function_calls_response_template = StringTemplate(
            """<function_calls_result>
  <name>{{name}}</name>
  <result>{{result}}</result>
</function_calls_result>""",
            template_format="jinja2",
        )
        self.filter_function_list = filter_function_list
        self.max_iterations = max_iterations
        self.max_execution_time: Optional[float] = None
        self.prefix_prompt = prefix_prompt
        if resources:
            self.resources = resources
        else:
            self.resources = Resources()

        if agent_name:
            self.agent_name = agent_name
        else:
            self.agent_name = generate_unique_id("react_agent")

        self.tools_name = []
        if self.tool_manager:
            for tool in self.tool_manager.tools:
                if self.filter_function_list and tool in self.filter_function_list:
                    self.tools_name.append(tool.name)
                else:
                    self.tools_name.append(tool.name)

        if self.mcp_client:
            for tool_name, _ in self.mcp_client.tools.items():
                if self.filter_function_list and tool_name in self.filter_function_list:
                    self.tools_name.append(tool_name)
                else:
                    self.tools_name.append(tool_name)

        self.curr_message: MessagesSet() = MessagesSet()

        self.is_interrupt = is_interrupt

    def get_llm(self) -> BaseAIChat:
        return self.llm

    def _build_system_prompt(self) -> MessagesSet:
        prefix_prompt = self.prefix_prompt + "\n\n" if self.prefix_prompt else ""

        _func_tool_descriptions = (
            self.tool_manager.tool_descriptions if self.tool_manager else "\n"
        )

        _mcp_tool_descriptions = (
            self.mcp_client.tool_descriptions if self.mcp_client else "\n"
        )

        return MessagesSet().add_system(
            prefix_prompt
            + "\n"
            + self.system_prompt_template.format(
                tool_descriptions=_func_tool_descriptions + _mcp_tool_descriptions
            )
        )

    def retry_llm_call(self, msg: MessagesSet):
        error_msg = ""
        logging.debug(f"react llm call start: {msg.get_format_messages()}")
        # print(msg.get_format_messages())
        for i in range(3):
            try:
                llm_resp = stream_to_string(self.llm.chat_for_stream(messages=msg))
                logging.debug(f"react llm call success: {llm_resp}")
                return llm_resp
            except Exception as e:
                logging.error(f"llm call error: {str(e)}")
                error_msg = str(e)
                time.sleep(1)
        else:
            raise ObserverException(error_msg)

    def _run(
        self,
        instruction: MessagesSet | str,
        verbose=True,
        **kwargs,
    ) -> Iterable[Union[ReactAction, ReactObservation, ReactResult]]:
        _system_prompt = self._build_system_prompt()
        self.curr_message = MessagesSet()
        if isinstance(instruction, str):
            self.curr_message.add_user(instruction)
        else:
            self.curr_message = instruction

        iterations = 0
        used_time = 0.0
        start_time = time.time()

        while self._should_continue(iterations, used_time):
            cell_id = generate_unique_id("react")
            _prompt = deepcopy(_system_prompt) + self.curr_message
            llm_resp = self.retry_llm_call(_prompt)

            self.curr_message.add_assistant(f"{remove_think(llm_resp)}")
            try:
                action_resp: ActionResponse = self._parse_llm_response(llm_resp)
            except ObserverException as e:
                logging.warning(f"ReactAction parse llm response error: {str(e)} ")
                self.curr_message.add_user(
                    self.function_calls_response_template.format(
                        name="unknown",
                        result=f"parse llm response error: {str(e)} ;response data : {remove_think(llm_resp)}",
                    )
                )
                continue

            _call_id = generate_unique_id("react")
            if verbose:
                if "final_answer" == action_resp.action.name:
                    pass
                else:
                    yield ReactAction(
                        **{
                            "action": {
                                "analysis": action_resp.analysis,
                                "action_name": action_resp.action.name,
                                "action_parameters": action_resp.action.args,
                            },
                            "cell_id": cell_id,
                            "call_id": _call_id,
                            "agent_name": self.agent_name,
                        }
                    )

            if "final_answer" == action_resp.action.name:
                _res = ReactResult(
                    analysis=action_resp.analysis,
                    result=action_resp.action.args["content"],
                    cell_id=cell_id,
                    agent_name=self.agent_name,
                )
                yield _res
                return

            try:
                tool_result = self.run_tool(
                    action_resp.action.name, action_resp.action.args
                )
                tool_exec_status = True

                # 如果是生成器，说明是子Agent (或者流式工具)，需要消费并透传
                if isinstance(tool_result, (GeneratorType, Iterator)):
                    _final_res = ""
                    for item in tool_result:
                        # 透传中间过程
                        yield item
                        # 捕获结果但不直接结束当前Agent loop
                        if isinstance(item, ReactResult):
                            _final_res = item.result
                    tool_result = _final_res

            except ObserverException as e:
                tool_result = f"call error: {str(e)}"
                tool_exec_status = False
            except Exception as e:
                tool_result = f"call error: {str(e)}"
                tool_exec_status = False
                logging.error(f"Tool execution error: {traceback.format_exc()}")

            self.curr_message.add_user(
                self.function_calls_response_template.format(
                    name=action_resp.action.name,
                    result=(
                        str(tool_result)
                        if isinstance(tool_result, (str, int, float))
                        else tool_result.to_str()
                    ),
                )
            )

            if verbose:
                yield ReactObservation(
                    **{
                        "observation": {
                            "result": (
                                str(tool_result)
                                if isinstance(tool_result, (str, int, float))
                                else tool_result.to_str()
                            ),
                            "view": (
                                tool_result
                                if isinstance(tool_result, (str, int, float)) or tool_result is None
                                else tool_result.to_view()
                            ),
                        },
                        "action": {
                            "analysis": action_resp.analysis,
                            "action_name": action_resp.action.name,
                            "action_parameters": action_resp.action.args,
                        },
                        "exec_status": tool_exec_status,
                        "cell_id": cell_id,
                        "call_id": _call_id,
                        "agent_name": self.agent_name,
                    }
                )
            iterations += 1
            used_time = time.time() - start_time

    def run_tool(self, tool_name: str, parameters: Union[str, dict]):
        if tool_name not in self.tools_name:
            return f"未知的工具调用 {tool_name}"

        if self.tool_manager and self.tool_manager.get_tool(tool_name):
            return self.tool_manager.run_tool(tool_name, parameters)
        elif self.mcp_client.get_tool_info(tool_name):
            return self.mcp_client.run_tool(tool_name, parameters)
        else:
            return f"未知的工具调用{tool_name}"

    def _should_continue(self, current_iteration: int, current_time_elapsed) -> bool:
        if self.max_iterations and current_iteration >= self.max_iterations:
            return False
        if self.max_execution_time and current_time_elapsed >= self.max_execution_time:
            return False
        if self.is_interrupt and self.is_interrupt():
            return False
        return True

    @staticmethod
    def _extract_response_block(xml_content: str) -> Optional[str]:
        """
        从XML内容中提取第一个完整的<response>...</response>块

        参数:
            xml_content: 包含可能多个XML块的字符串内容

        返回:
            包含完整<response>块的字符串，如果未找到则返回None
        """
        pattern = r"<response>.*?</response>"
        matches = re.findall(pattern, xml_content, re.DOTALL)

        if matches:
            return matches[0]
        return None

    @staticmethod
    def _extract_analysis(response_block: str) -> Optional[str]:
        """
        从<response>块中提取<analysis>标签的内容

        参数:
            response_block: 完整的<response> XML块字符串

        返回:
            <analysis>标签内的文本内容（如存在），如果未找到则返回None
        """
        try:
            # 方法1：使用lxml解析（推荐，更健壮）
            root = etree.fromstring(response_block)
            analysis = root.find(".//analysis")
            if analysis is not None:
                return analysis.text
            return None
        except:
            # 方法2：正则解析（备选方案）
            pattern = r"<analysis>(.*?)</analysis>"
            match = re.search(pattern, response_block, re.DOTALL)
            if match:
                return match.group(1).strip()
            return None

    @staticmethod
    def _parse_llm_response(llm_resp: str) -> ActionResponse:
        res = ReactAgent._extract_response_block(llm_resp)
        if not res:
            raise ObserverException("未找到完整的<response>...</response>块")

        analysis = ReactAgent._extract_analysis(res)
        if not analysis:
            raise ObserverException("未找到<analysis>...</analysis>标签")

        try:
            actions = parse_xml_tool_calls(res)
            for i in actions:
                if i.function_name == "response":
                    raise ObserverException(f"工具调用格式异常 {res}")
                return ActionResponse(
                    analysis=analysis,
                    action=ActionResponseArgs(
                        name=i.function_name,
                        args=i.parameters,
                    ),
                )
        except Exception as e:
            raise ObserverException(f"解析<response>...</response>块失败: {str(e)}")
