"""Personal memory retriever agent for retrieving personal memories through vector search."""

from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.op import BaseTool
from ....core.schema import Message
from ....core.utils import format_messages


class PersonalRetriever(BaseMemoryAgent):
    """Retrieve personal memories through vector search and history reading."""

    memory_type: MemoryType = MemoryType.PERSONAL

    async def build_messages(self) -> list[Message]:
        if self.context.get("query"):
            context = self.context.query
        elif self.context.get("messages"):
            context = self.description + "\n" + format_messages(self.context.messages)
        else:
            raise ValueError("input must have either `query` or `messages`")

        read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles")
        if read_all_profiles_tool is not None:
            all_profiles = await read_all_profiles_tool.call(
                memory_target=self.memory_target,
                service_context=self.service_context,
            )
        else:
            all_profiles = ""

        return [
            Message(
                role=Role.SYSTEM,
                content=self.prompt_format(
                    prompt_name="system_prompt",
                    memory_type=self.memory_type.value,
                    memory_target=self.memory_target,
                    user_profile=all_profiles,
                    context=context.strip(),
                ),
            ),
            Message(
                role=Role.USER,
                content=self.get_prompt("user_message"),
            ),
        ]

    async def _acting_step(
        self,
        assistant_message: Message,
        tools: list[BaseTool],
        step: int,
        stage: str = "",
        **kwargs,
    ) -> tuple[list[BaseTool], list[Message]]:
        """Execute tool calls with memory context."""
        return await super()._acting_step(
            assistant_message,
            tools,
            step,
            memory_type=self.memory_type.value,
            memory_target=self.memory_target,
            retrieved_nodes=self.retrieved_nodes,
            **kwargs,
        )

    async def execute(self):
        result = await super().execute()
        answer = result["answer"]
        if "MEMORY_NOT_FOUND" in answer:
            result["answer"] = "\n".join(
                [
                    n.format(
                        include_memory_id=False,
                        include_when_to_use=False,
                        include_content=True,
                        include_message_time=False,
                        ref_memory_id_key="",
                    )
                    for n in self.retrieved_nodes
                ],
            )

        result["retrieved_nodes"] = self.retrieved_nodes
        return result
