# -*- coding: utf-8 -*-
import asyncio
import json
import time
import traceback
from typing import Any, Optional, Type

from loguru import logger
from pydantic import BaseModel

from agentscope.agent import ReActAgent
from agentscope.model import ChatModelBase
from agentscope.formatter import FormatterBase
from agentscope.memory import MemoryBase
from agentscope.tracing import trace_reply
from agentscope.message import Msg, TextBlock, ToolUseBlock, ToolResultBlock

from alias.agent.tools import AliasToolkit
from alias.agent.utils.constants import DEFAULT_PLANNER_NAME
from alias.agent.utils.agent_save_state import AliasAgentStates
from alias.agent.utils.constants import MODEL_MAX_RETRIES


class AliasAgentBase(ReActAgent):
    def __init__(
        self,
        name: str,
        model: ChatModelBase,
        formatter: FormatterBase,
        memory: MemoryBase,
        toolkit: AliasToolkit,
        session_service: Any,
        state_saving_dir: Optional[str] = None,
        sys_prompt: Optional[str] = None,
        max_iters: int = 10,
    ):
        super().__init__(
            name=name,
            sys_prompt=sys_prompt,
            model=model,
            formatter=formatter,
            memory=memory,
            toolkit=toolkit,
            max_iters=max_iters,
        )

        self.session_service = session_service
        self.message_sending_mapping = {}
        self.state_saving_dir = state_saving_dir

    async def _reasoning(self):
        """Override _reasoning to add retry logic."""

        # Call the parent class's _reasoning method directly to
        # avoid double hook execution
        # We need to call the underlying implementation without hooks
        async def call_parent_reasoning():
            # Get the original method from the parent class before
            # metaclass processing
            # Access the method from the class that defines it
            # (before metaclass wrapping)
            original_method = ReActAgent.__dict__["_reasoning"]
            # Check if this is the wrapped version by looking for
            # the wrapper attributes
            if hasattr(original_method, "__wrapped__"):
                # This is the wrapped version, get the original
                original_method = original_method.__wrapped__
            return await original_method(self)

        for i in range(MODEL_MAX_RETRIES - 1):
            try:
                return await call_parent_reasoning()
            except Exception:
                logger.warning(
                    f"Reasoning fail at attempt {i + 1}. "
                    f"Max attempts {MODEL_MAX_RETRIES}\n"
                    f"{traceback.format_exc()}",
                )
                memory_msgs = await self.memory.get_memory()
                mem_len = len(memory_msgs)
                # ensure the last message has no tool_use before next attempt
                if mem_len > 0 and memory_msgs[-1].has_content_blocks(
                    "tool_use",
                ):
                    await self.memory.delete(index=mem_len - 1)
                time.sleep(2)

        # final attempt
        await call_parent_reasoning()

    @trace_reply
    async def reply(
        self,
        msg: Msg | list[Msg] | None = None,
        structured_model: Type[BaseModel] | None = None,
    ) -> Msg:
        """Generate a reply based on the current state and input arguments.

        TODO: (part 1)
        this is just a monkey patch for AS when not support interruption
        during tool call; to be remove when AS framework updated

        Args:
            msg (`Msg | list[Msg] | None`, optional):
                The input message(s) to the agent.
            structured_model (`Type[BaseModel] | None`, optional):
                The required structured output model. If provided, the agent
                is expected to generate structured output in the `metadata`
                field of the output message.

        Returns:
            `Msg`:
                The output message generated by the agent.
        """
        await self.memory.add(msg)

        # Long-term memory retrieval
        if self._static_control:
            # Retrieve information from the long-term memory if available
            retrieved_info = await self.long_term_memory.retrieve(msg)
            if retrieved_info:
                await self.memory.add(
                    Msg(
                        name="long_term_memory",
                        content="<long_term_memory>The content below are "
                        "retrieved from long-term memory, which maybe "
                        f"useful:\n{retrieved_info}"
                        f"</long_term_memory>",
                        role="user",
                    ),
                )

        self._required_structured_model = structured_model
        # Record structured output model if provided
        if structured_model:
            self.toolkit.set_extended_model(
                self.finish_function_name,
                structured_model,
            )

        # The reasoning-acting loop
        reply_msg = None
        for _ in range(self.max_iters):
            msg_reasoning = await self._reasoning()

            futures = [
                self._acting(tool_call)
                for tool_call in msg_reasoning.get_content_blocks(
                    "tool_use",
                )
            ]

            # Parallel tool calls or not
            if self.parallel_tool_calls:
                acting_responses = await asyncio.gather(*futures)

            else:
                # Sequential tool calls
                acting_responses = [await _ for _ in futures]

            # Find the first non-None replying message from the acting
            for acting_msg in acting_responses:
                reply_msg = reply_msg or acting_msg
                # TODO: monkey patch happens here
                if (
                    isinstance(reply_msg, Msg)
                    and reply_msg.metadata
                    and reply_msg.metadata.get("is_interrupted", False)
                ):
                    raise asyncio.CancelledError()

            if reply_msg:
                break

        # When the maximum iterations are reached
        if reply_msg is None:
            reply_msg = await self._summarizing()

        # Post-process the memory, long-term memory
        if self._static_control:
            await self.long_term_memory.record(
                [
                    *([*msg] if isinstance(msg, list) else [msg]),
                    *await self.memory.get_memory(),
                    reply_msg,
                ],
            )

        await self.memory.add(reply_msg)
        return reply_msg

    async def _acting(self, tool_call: ToolUseBlock) -> Msg | None:
        """Perform the acting process.

        TODO: (part 2)
        this is just a monkey patch for AS when not support interruption
        during tool call; can be remove when AS framework updated

        Args:
            tool_call (`ToolUseBlock`):
                The tool use block to be executed.

        Returns:
            `Union[Msg, None]`:
                Return a message to the user if the `_finish_function` is
                called, otherwise return `None`.
        """

        tool_res_msg = Msg(
            "system",
            [
                ToolResultBlock(
                    type="tool_result",
                    id=tool_call["id"],
                    name=tool_call["name"],
                    output=[],
                ),
            ],
            "system",
        )
        try:
            # Execute the tool call
            tool_res = await self.toolkit.call_tool_function(tool_call)

            response_msg = None
            # Async generator handling
            async for chunk in tool_res:
                # Turn into a tool result block
                tool_res_msg.content[0][  # type: ignore[index]
                    "output"
                ] = chunk.content

                # todo: monkey patch to pass the metadata
                if chunk.metadata:
                    if tool_res_msg.metadata is None:
                        tool_res_msg.metadata = {}
                    for key, value in chunk.metadata.items():
                        try:
                            # verify it's JSON-serializable
                            json.dumps(value)
                            tool_res_msg.metadata[key] = value
                        except (TypeError, ValueError):
                            # Skip non-serializable values
                            pass

                # Skip the printing of the finish function call
                if (
                    tool_call["name"] != self.finish_function_name
                    or tool_call["name"] == self.finish_function_name
                    and not chunk.metadata.get("success")
                ):
                    await self.print(tool_res_msg, chunk.is_last)

                # Return message if generate_response is called successfully
                if tool_call[
                    "name"
                ] == self.finish_function_name and chunk.metadata.get(
                    "success",
                    True,
                ):
                    response_msg = chunk.metadata.get("response_msg")
                elif chunk.is_interrupted:
                    # TODO: monkey patch happens here
                    response_msg = tool_res_msg
                    if response_msg.metadata is None:
                        response_msg.metadata = {"is_interrupted": True}
                    else:
                        response_msg.metadata["is_interrupted"] = True

            return response_msg
        finally:
            # Record the tool result message in the memory
            await self.memory.add(tool_res_msg)

    async def handle_interrupt(
        self,
        _msg: Msg | list[Msg] | None = None,
    ) -> Msg:
        """
        The post-processing logic when the reply is interrupted by the
        user or something else.
        """
        response_msg = Msg(
            self.name,
            content=[
                TextBlock(
                    type="text",
                    text="I got interrupted by the user. "
                    "Pivot to handle the user's new request.",
                ),
            ],
            role="assistant",
            metadata={},
        )
        await self.memory.add(response_msg)

        # update and save agent states
        global_state = await self.session_service.get_state()
        if global_state is None:
            global_state = AliasAgentStates()
        else:
            global_state = AliasAgentStates(**global_state)
        global_state.agent_states[self.name] = self.state_dict()
        await self.session_service.create_state(
            content=global_state.model_dump(),
        )

        if self.name == DEFAULT_PLANNER_NAME:
            return response_msg
        else:
            raise asyncio.CancelledError
