"""Auto plan chat manager agent."""

import logging
from typing import Dict, List, Optional, Tuple

from dbgpt.core.interface.message import ModelMessageRoleType

from ..action.base import ActionOutput
from ..agent import Agent, AgentMessage
from ..agent_manage import mentioned_agents, participant_roles
from ..base_agent import ConversableAgent
from ..base_team import ManagerAgent
from ..memory.gpts.base import GptsPlan
from ..plan.planner_agent import PlannerAgent
from ..profile import DynConfig, ProfileConfig
from ..schema import Status

logger = logging.getLogger(__name__)


class AutoPlanChatManager(ManagerAgent):
    """A chat manager agent that can manage a team chat of multiple agents."""

    profile: ProfileConfig = ProfileConfig(
        name=DynConfig(
            "AutoPlanChatManager",
            category="agent",
            key="dbgpt_agent_plan_team_auto_plan_profile_name",
        ),
        role=DynConfig(
            "PlanManager",
            category="agent",
            key="dbgpt_agent_plan_team_auto_plan_profile_role",
        ),
        goal=DynConfig(
            "Advance the task plan generated by the planning agent. If the plan "
            "does not pre-allocate an agent, it needs to be coordinated with the "
            "appropriate agent to complete.",
            category="agent",
            key="dbgpt_agent_plan_team_auto_plan_profile_goal",
        ),
        desc=DynConfig(
            "Advance the task plan generated by the planning agent.",
            category="agent",
            key="dbgpt_agent_plan_team_auto_plan_profile_desc",
        ),
    )

    def __init__(self, **kwargs):
        """Create a new AutoPlanChatManager instance."""
        super().__init__(**kwargs)

    async def process_rely_message(
        self, conv_id: str, now_plan: GptsPlan, speaker: Agent
    ):
        """Process the dependent message."""
        rely_prompt = None
        rely_messages: List[Dict] = []

        if now_plan.rely and len(now_plan.rely) > 0:
            rely_tasks_list = now_plan.rely.split(",")
            rely_tasks_list_int = [int(i) for i in rely_tasks_list]
            rely_tasks = self.memory.plans_memory.get_by_conv_id_and_num(
                conv_id, rely_tasks_list_int
            )
            if rely_tasks:
                rely_prompt = (
                    "Read the result data of the dependent steps in the above"
                    " historical message to complete the current goal:"
                )
                for rely_task in rely_tasks:
                    rely_messages.append(
                        {
                            "content": rely_task.sub_task_content,
                            "role": ModelMessageRoleType.HUMAN,
                            "name": rely_task.sub_task_agent,
                        }
                    )
                    rely_messages.append(
                        {
                            "content": rely_task.result,
                            "role": ModelMessageRoleType.AI,
                            "name": rely_task.sub_task_agent,
                        }
                    )
        return rely_prompt, rely_messages

    def select_speaker_msg(self, agents: List[Agent]) -> str:
        """Return the message for selecting the next speaker."""
        agent_names = [agent.name for agent in agents]
        return (
            "You are in a role play game. The following roles are available:\n"
            f"   {participant_roles(agents)}.\n"
            "   Read the following conversation.\n"
            f"   Then select the next role from {agent_names} to play.\n"
            "   The role can be selected repeatedly.Only return the role."
        )

    async def select_speaker(
        self,
        last_speaker: Agent,
        selector: Agent,
        now_goal_context: Optional[str] = None,
        pre_allocated: Optional[str] = None,
    ) -> Tuple[Agent, Optional[str]]:
        """Select the next speaker."""
        agents = self.agents

        if pre_allocated:
            # Preselect speakers
            logger.info(f"Preselect speakers:{pre_allocated}")
            name = pre_allocated
            model = None
        else:
            # auto speaker selection
            # TODO selector a_thinking It has been overwritten and cannot be used.
            agent_names = [agent.name for agent in agents]
            fina_name, model = await selector.thinking(
                messages=[
                    AgentMessage(
                        role=ModelMessageRoleType.HUMAN,
                        content="Read and understand the following task content and"
                        " assign the appropriate role to complete the task.\n"
                        f"Task content: {now_goal_context},\n"
                        f"Select the role from: {agent_names},\n"
                        f"Please only return the role, such as: {agents[0].name}",
                    )
                ],
                prompt=self.select_speaker_msg(agents),
            )
            if not fina_name:
                raise ValueError("Unable to select next speaker!")
            else:
                name = fina_name

        # If exactly one agent is mentioned, use it. Otherwise, leave the OAI response
        # unmodified
        mentions = mentioned_agents(name, agents)
        if len(mentions) == 1:
            name = next(iter(mentions))
        else:
            logger.warning(
                "GroupChat select_speaker failed to resolve the next speaker's name. "
                f"This is because the speaker selection OAI call returned:\n{name}"
            )

        # Return the result
        try:
            return self.agent_by_name(name), model
        except Exception as e:
            logger.exception(f"auto select speaker failed!{str(e)}")
            raise ValueError("Unable to select next speaker!")

    async def act(
        self,
        message: AgentMessage,
        sender: Agent,
        reviewer: Optional[Agent] = None,
        is_retry_chat: bool = False,
        last_speaker_name: Optional[str] = None,
        **kwargs,
    ) -> ActionOutput:
        """Perform an action based on the received message."""
        if not sender:
            return ActionOutput(
                is_exe_success=False,
                content="The sender cannot be empty!",
            )
        speaker: Agent = sender
        final_message = message.content
        rounds = message.rounds
        for i in range(self.max_round):
            if not self.memory:
                return ActionOutput(
                    is_exe_success=False,
                    content="The memory cannot be empty!",
                )
            plans = self.memory.plans_memory.get_by_conv_id(
                self.not_null_agent_context.conv_id
            )

            if not plans or len(plans) <= 0:
                if i > 3:
                    return ActionOutput(
                        is_exe_success=False,
                        content="Retrying 3 times based on current application "
                        "resources still fails to build a valid plan！",
                    )
                planner: ConversableAgent = (
                    await PlannerAgent()
                    .bind(self.memory)
                    .bind(self.agent_context)
                    .bind(self.llm_config)
                    .bind_agents(self.agents)
                    .build()
                )

                plan_message = await planner.generate_reply(
                    received_message=AgentMessage.from_llm_message(
                        {"content": message.content, "rounds": rounds}
                    ),
                    sender=self,
                    reviewer=reviewer,
                )
                rounds = plan_message.rounds
                await planner.send(
                    message=plan_message, recipient=self, request_reply=False
                )
            else:
                todo_plans = [
                    plan
                    for plan in plans
                    if plan.state in [Status.TODO.value, Status.RETRYING.value]
                ]
                if not todo_plans or len(todo_plans) <= 0:
                    # The plan has been fully executed and a success message is sent
                    # to the user.
                    # complete
                    return ActionOutput(
                        is_exe_success=True,
                        content=final_message,  # work results message
                    )
                else:
                    try:
                        now_plan: GptsPlan = todo_plans[0]
                        current_goal_message = AgentMessage(
                            content=now_plan.sub_task_content,
                            current_goal=now_plan.sub_task_content,
                            context={
                                "plan_task": now_plan.sub_task_content,
                                "plan_task_num": now_plan.sub_task_num,
                            },
                            rounds=rounds + 1,
                        )
                        # select the next speaker
                        speaker, model = await self.select_speaker(
                            speaker,
                            self,
                            now_plan.sub_task_content,
                            now_plan.sub_task_agent,
                        )
                        # Tell the speaker the dependent history information
                        rely_prompt, rely_messages = await self.process_rely_message(
                            conv_id=self.not_null_agent_context.conv_id,
                            now_plan=now_plan,
                            speaker=speaker,
                        )
                        if rely_prompt:
                            current_goal_message.content = (
                                rely_prompt + current_goal_message.content
                            )

                        await self.send(
                            message=current_goal_message,
                            recipient=speaker,
                            reviewer=reviewer,
                            request_reply=False,
                        )
                        agent_reply_message = await speaker.generate_reply(
                            received_message=current_goal_message,
                            sender=self,
                            reviewer=reviewer,
                            rely_messages=AgentMessage.from_messages(rely_messages),
                        )
                        is_success = agent_reply_message.success
                        reply_message = agent_reply_message.to_llm_message()
                        await speaker.send(
                            agent_reply_message, self, reviewer, request_reply=False
                        )
                        rounds = agent_reply_message.rounds

                        plan_result = ""
                        final_message = reply_message["content"]
                        if is_success:
                            if reply_message:
                                action_report = agent_reply_message.action_report
                                if action_report:
                                    plan_result = action_report.content
                                    final_message = action_report.view

                            # The current planned Agent generation verification is
                            # successful
                            # Plan executed successfully
                            self.memory.plans_memory.complete_task(
                                self.not_null_agent_context.conv_id,
                                now_plan.sub_task_num,
                                plan_result,
                            )
                        else:
                            plan_result = reply_message["content"]
                            self.memory.plans_memory.update_task(
                                self.not_null_agent_context.conv_id,
                                now_plan.sub_task_num,
                                Status.FAILED.value,
                                now_plan.retry_times + 1,
                                speaker.name,
                                "",
                                plan_result,
                            )
                            return ActionOutput(
                                is_exe_success=False, content=plan_result
                            )

                    except Exception as e:
                        logger.exception(
                            f"An exception was encountered during the execution of the"
                            f" current plan step.{str(e)}"
                        )
                        return ActionOutput(
                            is_exe_success=False,
                            content=f"An exception was encountered during the execution"
                            f" of the current plan step.{str(e)}",
                        )
        return ActionOutput(
            is_exe_success=False,
            content=f"Maximum number of dialogue rounds exceeded.{self.max_round}",
        )
