from abc import ABC, abstractmethod

import contexttimer
from judy_core.agents.errors import AgentError
from judy_core.agents.ids import NOTIFICATION_CENTER_ACTOR_ID
from judy_core.agents.stateful_llm_agent import StatefulLLMAgent
from judy_core.events import DebugToolResult, TaskFailResult, ToolCall, ToolResult
from judy_core.messages import EnvironmentNotificationList, GetEnvironmentNotification
from judy_core.utils.messages import convert_content_to_openai_format
from judy_core.utils.tool import ToolRegistry
from litellm.types.completion import (
    ChatCompletionToolMessageParam,
    ChatCompletionUserMessageParam,
)
from tongcell.core import MessageContext, message_handler


class ReactAgent(StatefulLLMAgent, ABC):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self._internal_tool_registry = ToolRegistry(tools=[])

        self._step_count = 0
        self._tool_call_count = 0

        self._tool_calls = []
        self._tool_results = dict[str, ToolResult]()

    @abstractmethod
    def _react_exit_prompt(self) -> str: ...

    @abstractmethod
    async def _on_turn_done(self, ctx: MessageContext) -> None: ...

    async def _do_turn(self, ctx: MessageContext) -> None:
        for i in range(0, 5):
            ok = await self._do_turn_inner(ctx)
            if ok:
                return
        await self._publish_output_event(
            TaskFailResult(content="executor stuck in do_turn loop")
        )

    async def _do_turn_inner(self, ctx: MessageContext) -> bool:
        await self._add_notification_message(ctx)

        choice = await self._acompletion(
            messages=self._messages,
            tools=self._tools,
            cache_control_injection_points=[
                {
                    "location": "message",
                    "index": -1,
                }
            ],
        )
        if choice.message.tool_calls:
            self._tool_message = choice.message
            self._tool_calls = choice.message.tool_calls
            self._tool_results.clear()
            internal_tool_names = self._internal_tool_registry.get_tool_names()
            for tool_call in choice.message.tool_calls:
                tool_call_id = tool_call.id
                function_name = tool_call.function.name
                function_arguments = tool_call.function.arguments
                if not function_name:
                    raise AgentError("empty function name")
                if function_name in internal_tool_names:
                    with contexttimer.Timer() as t:
                        tool_exec_result = await self._internal_tool_registry.call_tool(
                            function_name, function_arguments
                        )
                    await self._publish_output_event(
                        DebugToolResult(
                            source=str(self._id),
                            tool_call_id=tool_call_id,
                            name=function_name,
                            content=tool_exec_result.content,
                            time_cost=t.elapsed,
                        )
                    )
                    await ctx.send(
                        ToolResult(
                            id=tool_call_id,
                            name=function_name,
                            content=tool_exec_result.content,
                        ),
                        self._id,
                    )
                else:
                    await self._publish_output_event(
                        ToolCall(
                            source=str(self._id),
                            tool_call_id=tool_call_id,
                            name=function_name,
                            arguments=function_arguments,
                        ),
                    )
            return True
        else:
            await self._debug_log(text="did not issue any tool calls during react step")
            self._messages.append(
                {
                    "role": "user",
                    "content": self._react_exit_prompt(),
                }
            )
            return False

    async def _add_notification_message(self, ctx: MessageContext):
        notification_list = await ctx.request(
            GetEnvironmentNotification(),
            NOTIFICATION_CENTER_ACTOR_ID,
            response_type=EnvironmentNotificationList,
        )
        if not notification_list.notifications:
            return
        text_segs = ["There are some notifications for you. "]
        for notification in notification_list.notifications:
            text_segs.extend(
                [
                    "<notification>",
                    notification.text,
                    "</notification>",
                ]
            )
        message = ChatCompletionUserMessageParam(
            role="user", content="\n".join(text_segs)
        )
        await self._publish_model_input([message])
        self._messages.append(message)

    @message_handler()
    async def receive_tool_result(
        self, tool_result: ToolResult, ctx: MessageContext
    ) -> None:
        id_found = False
        for tool_call in self._tool_calls:
            if tool_call.id == tool_result.id:
                id_found = True
                break
        if not id_found:
            raise AgentError(f"tool_call id not found: {tool_result.id}")
        if self._tool_results.get(tool_result.id) is not None:
            raise AgentError(f"tool_call result already set: {tool_result.id}")
        self._tool_results[tool_result.id] = tool_result
        self._tool_call_count += 1

        if len(self._tool_calls) > len(self._tool_results):
            return None

        tool_turn_messages: list = [self._tool_message]
        media_content_parts = []
        for tool_call in self._tool_calls:
            temp_tool_result = self._tool_results[tool_call.id]
            tool_result_content = convert_content_to_openai_format(
                temp_tool_result.content, self._model_config.model
            )

            if (
                "claude" in self._model_config.model
                and "openrouter" in self._model_config.model
            ) or "gemini" in self._model_config.model:
                # openroute上的claude api会过调掉工具返回的多模态内容。
                # gemini 3.0虽然支持工具返回多模态内容，但litellm在转换时将工具返回的内容转换为了多个part，没有转换成gemini 3.0需要的functionResponse part的子part。
                # 导致gemini经常出现乱码，message中出现reasoning内容，不断重复相同文本内容/tool call，tool call内容非法等诡异问题。
                if not isinstance(tool_result_content, str):
                    tool_result_text_content = []
                    for part in tool_result_content:
                        if part.get("type") == "text":
                            tool_result_text_content.append(part)
                        else:
                            media_content_parts.append(
                                {
                                    "type": "text",
                                    "text": "The following content is the media file read by tool call `{} {}`".format(
                                        tool_call.function.name,
                                        tool_call.function.arguments,
                                    ),
                                }
                            )
                            media_content_parts.append(part)
                    tool_result_content = tool_result_text_content

            tool_turn_messages.append(
                ChatCompletionToolMessageParam(
                    tool_call_id=tool_call.id,
                    content=tool_result_content,
                    role="tool",
                )
            )

        self._messages.extend(tool_turn_messages)
        if media_content_parts:
            self._messages.append(
                ChatCompletionUserMessageParam(
                    role="user",
                    content=media_content_parts,
                )
            )

        await self._on_turn_done(ctx)
