"""Tool memory retrieval operation module.

This module provides functionality to retrieve tool memories from a vector store
based on tool names, format them into structured documents, and match them with
the requested tools.
"""

from typing import List

from flowllm.core.context import C
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import VectorNode
from loguru import logger

from reme_ai.schema.memory import ToolMemory, vector_node_to_memory


@C.register_op()
class RetrieveToolMemoryOp(BaseAsyncOp):
    """Retrieves tool memories from vector store based on tool names.

    This operation searches for tool memories in the vector store using tool names,
    validates that the retrieved memories match the requested tools, and formats
    them into a structured document format for use in the context.
    """

    file_path: str = __file__

    @staticmethod
    def _format_tool_memories(memories: List[ToolMemory]) -> str:
        """Format tool memories into a structured document format.

        Args:
            memories: List of ToolMemory objects to format.

        Returns:
            A formatted string containing all tool memories with separators.
        """
        lines = [f"Retrieved {len(memories)} tool memory(ies):\n"]

        for idx, memory in enumerate(memories, 1):
            lines.append(f"Tool: {memory.when_to_use}")
            lines.append(memory.content)

            if idx < len(memories):
                lines.append("\n---\n")

        return "\n".join(lines)

    async def async_execute(self):
        """Execute the tool memory retrieval operation.

        This method:
        1. Extracts tool names from context
        2. Searches for each tool in the vector store
        3. Validates that retrieved memories match the requested tools
        4. Formats the memories into a structured document
        5. Stores the results in context response

        The operation expects 'tool_names' in the context, which should be a
        comma-separated string of tool names. For each tool name, it retrieves
        the top matching memory from the vector store and validates that it
        matches exactly.
        """
        tool_names: str = self.context.get("tool_names", "")
        workspace_id: str = self.context.workspace_id

        if not tool_names:
            logger.warning("tool_names is empty, skipping processing")
            self.context.response.answer = "tool_names is required"
            self.context.response.success = False
            return

        # Split tool names by comma
        tool_name_list = [name.strip() for name in tool_names.split(",") if name.strip()]
        logger.info(f"workspace_id={workspace_id} retrieving {len(tool_name_list)} tools: {tool_name_list}")

        # Search for each tool in the vector store
        matched_tool_memories: List[ToolMemory] = []

        for tool_name in tool_name_list:
            nodes: List[VectorNode] = await self.vector_store.async_search(
                query=tool_name,
                workspace_id=workspace_id,
                top_k=1,
            )

            if nodes:
                top_node = nodes[0]
                memory = vector_node_to_memory(top_node)

                # Ensure it's a ToolMemory and when_to_use matches
                if isinstance(memory, ToolMemory) and memory.when_to_use == tool_name:
                    matched_tool_memories.append(memory)
                    logger.info(
                        f"Found tool_memory for tool_name={tool_name}, "
                        f"memory_id={memory.memory_id}, "
                        f"total_calls={len(memory.tool_call_results)}",
                    )
                else:
                    logger.warning(f"No exact match found for tool_name={tool_name}")
            else:
                logger.warning(f"No memory found for tool_name={tool_name}")

        if not matched_tool_memories:
            logger.info("No matching tool memories found")
            self.context.response.answer = "No matching tool memories found"
            self.context.response.success = False
            return

        # Format tool memories as document
        formatted_answer = self._format_tool_memories(matched_tool_memories)

        # Set response
        self.context.response.answer = formatted_answer
        self.context.response.success = True
        self.context.response.metadata["memory_list"] = matched_tool_memories

        # Log retrieval results
        for memory in matched_tool_memories:
            logger.info(
                f"Retrieved tool: {memory.when_to_use}, "
                f"total_calls={len(memory.tool_call_results)}, "
                f"content_length={len(memory.content)}",
            )
