from copy import copy

from judy_core.agents.stateful_llm_agent import StatefulLLMAgent
from judy_core.agents.types import ModelInputMessage
from judy_core.messages import MemberAnswer, MemberQuestion, TeamMemberName
from tongcell.core import MessageContext, message_handler

MEMBER_QA_INSTRUCTION_TEMPLATE = """\
从现在开始你进入与团队成员({question_member_name})的问答流程，对方提问，你来回答(每条回复不超过8000字，超过后中间部分会被截断，开头和尾部会被保留)。
注意：**在问答流程中，你不能调用任何工具**。
"""


class QAAgent(StatefulLLMAgent):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._qa_conversations: dict[TeamMemberName, list[ModelInputMessage]] = {}

    @message_handler()
    async def answer_question(
        self, question: MemberQuestion, ctx: MessageContext
    ) -> MemberAnswer:
        question_member_name = ctx.envelope.sender.name
        qa_messages = self._qa_conversations.setdefault(
            question_member_name,
            [
                {
                    "role": "user",
                    "content": MEMBER_QA_INSTRUCTION_TEMPLATE.format(
                        question_member_name=question_member_name
                    ),
                }
            ],
        )
        qa_messages.append({"role": "user", "content": question.text})

        messages = copy(self._messages)
        messages.extend(qa_messages)
        choice = await self._acompletion_without_tool_calls(
            messages=messages,
            tools=self._tools,
            publish_model_output=False,
        )
        qa_messages.append(choice.message)

        assert choice.message.content is not None
        return MemberAnswer(text=choice.message.content)
