import traceback
from copy import copy
from typing import Any, Sequence, cast

import contexttimer
import litellm
from judy_core.agents.errors import AgentError, ModelBlankResponseError
from judy_core.agents.ids import CONFIGURATION_ACTOR_ID
from judy_core.events import (
    DebugModelInput,
    DebugModelOutput,
    DebugTaskLog,
    OutputEvent,
    TaskFailResult,
)
from judy_core.messages import GetModelConfig, GetTaskConfig, ModelConfig, TaskConfig
from judy_core.topics import publish_output_event
from judy_core.utils.llm import acompletion
from litellm.types.utils import Choices, Usage
from tongcell.core import (
    ActorStartMessage,
    MessageContext,
    RoutedActor,
    message_handler,
)
from tongcell.util import getLogger

logger = getLogger(__name__)


LLM_TOOLS_DISABLED_INSTRUCTION = """\
**Ignore this message. This message is only meant to remind you \
that you are now prohibited from calling any tools, and you should \
not give an empty response.**
"""


class LLMAgent(RoutedActor):
    @message_handler()
    async def on_start(self, msg: ActorStartMessage, ctx: MessageContext) -> None:
        self._model_config = await ctx.request(
            GetModelConfig(), CONFIGURATION_ACTOR_ID, response_type=ModelConfig
        )
        self._task_config = await ctx.request(
            GetTaskConfig(), CONFIGURATION_ACTOR_ID, response_type=TaskConfig
        )
        self._ctx = ctx

    async def on_handler_exception(
        self, message: Any, ex: Exception, ctx: MessageContext
    ) -> None:
        await self._publish_output_event(
            TaskFailResult(
                source=str(self.id), content="".join(traceback.format_exception(ex))
            )
        )
        return None

    def _need_explicit_cache(self) -> bool:
        need_explicit_cache = "claude" in self._model_config.model
        return need_explicit_cache

    async def _acompletion(
        self,
        messages,
        *,
        publish_model_output: bool = True,
        cache_control_injection_points: list | None = None,
        **kwargs,
    ) -> Choices:
        # logger.info(
        #     "messages %s", json.dumps(kwargs.get("messages"), ensure_ascii=False)
        # )
        if not self._need_explicit_cache():
            cache_control_injection_points = None
        with contexttimer.Timer() as t:
            litellm_completion_option = self._model_config.litellm_completion_option
            response = await acompletion(
                model=self._model_config.model,
                messages=messages,
                judy_option=self._model_config.judy_completion_option,
                judy_logger=self._debug_log,
                **litellm_completion_option.model_dump(),
                parallel_tool_calls=True,
                cache_control_injection_points=cache_control_injection_points,
                **kwargs,
            )
            if response is None:
                raise ModelBlankResponseError(
                    f"No choices returned from LLM: {response}"
                )

        choice = cast(Choices, response.choices[0])
        # TODO publish的消息可能还没来得及处理，runtime就出现异常把receiver stop了
        reasoning_content = getattr(choice.message, "reasoning_content", None)
        usage: Usage | None = getattr(response, "usage", None)

        cost_money = 0.0
        hidden_params = getattr(response, "_hidden_params", None)
        if isinstance(hidden_params, dict):
            response_cost = hidden_params.get("response_cost", 0.0)
            if response_cost:
                cost_money = response_cost

        if cost_money == 0.0:
            try:
                cost_money = litellm.completion_cost(completion_response=response)
            except Exception:
                pass

        if publish_model_output:
            await self._publish_output_event(
                DebugModelOutput(
                    source=str(self.id),
                    response_id=response.id,
                    model=response.model,
                    tool_calls=choice.message.tool_calls,
                    content=choice.message.content,
                    reasoning_content=reasoning_content,
                    usage=usage,
                    hidden_params=hidden_params,
                    cost_money=cost_money,
                    cost_time=t.elapsed,
                ),
            )

        return choice

    async def _acompletion_without_tool_calls(self, messages, **kwargs) -> Choices:
        messages = copy(messages)
        last_message = messages[-1]
        for i in range(0, 3):
            choice = await self._acompletion(messages, **kwargs)
            if choice.message.tool_calls or (not choice.message.content):
                await self._debug_log(
                    "agent {} try to call tools(or message.content is empty) while tool is disabled, choice:{}, last_message: {}".format(
                        self._id, choice, last_message
                    )
                )
                messages.append(
                    {
                        "role": "user",
                        "content": LLM_TOOLS_DISABLED_INSTRUCTION,
                    }
                )
                continue
            else:
                return choice
        error_content = "{} failed to generate a message without tool calls and text is not empty".format(
            self._id
        )
        await self._debug_log(error_content)
        raise AgentError(error_content)

    async def _publish_model_input(self, messages: Sequence) -> None:
        for message in messages:
            role = message.get("role")
            content = message.get("content")
            if isinstance(content, list):
                content = "".join([part.get("text", "") for part in content])
            if isinstance(role, str) and isinstance(content, str):
                await self._publish_output_event(
                    DebugModelInput(
                        source=str(self.id),
                        role=role,
                        content=content,
                    )
                )

    async def _publish_output_event(self, event: OutputEvent) -> None:
        await publish_output_event(self._ctx, event)

    async def _debug_log(self, text: str):
        await self._publish_output_event(
            DebugTaskLog(content="agent:{} {}".format(self._id, text))
        )
