"""
Tool for searching memories by entity relationships.
Allows agents to find all memories related to a specific entity (Agent, Technology, Project, etc.)
"""
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
from typing import Optional
from src.memory.graph_wrapper import GraphWrapper
import logging

logger = logging.getLogger("AntigravityEntitySearch")

class EntitySearchInput(BaseModel):
    entity_name: str = Field(..., description="Name of the entity to search for (e.g., 'Arthur Mendes', 'Zabbix', 'Infrastructure Project')")
    entity_type: Optional[str] = Field(
        default=None,
        description="Type of entity: Agent, Technology, Project, Problem, Solution. If not specified, searches all types."
    )
    limit: Optional[int] = Field(
        default=10,
        description="Maximum number of memories to return"
    )

class SearchByEntityTool(BaseTool):
    name: str = "Search Memories by Entity"
    description: str = (
        "Search for all memories related to a specific entity (person, technology, project, etc.) "
        "using the relationship graph. This is useful when you need to find everything related to "
        "a specific agent, technology, or project, not just semantically similar content. "
        "Example: 'Find all memories about Zabbix' or 'Show me everything related to Arthur Mendes'."
    )
    args_schema: type = EntitySearchInput

    def _run(self, entity_name: str, entity_type: Optional[str] = None, limit: int = 10) -> str:
        """
        Busca memórias relacionadas a uma entidade específica usando o grafo de relacionamentos.
        
        Args:
            entity_name: Nome da entidade
            entity_type: Tipo da entidade (Agent, Technology, Project, etc.)
            limit: Limite de resultados
        """
        try:
            memories = GraphWrapper.find_memories_by_entity(
                entity_name=entity_name,
                entity_type=entity_type,
                limit=limit
            )
            
            if not memories:
                return f"No memories found related to entity '{entity_name}'" + (f" of type '{entity_type}'" if entity_type else "") + "."
            
            formatted = [f"Found {len(memories)} memories related to '{entity_name}':\n"]
            
            for i, mem in enumerate(memories, 1):
                preview = mem.get('preview', 'N/A')
                mem_type = mem.get('type', 'unknown')
                entity_type_display = mem.get('entity_type', 'N/A')
                related_count = mem.get('related_count', 0)
                
                formatted.append(f"{i}. [{mem_type}] {preview}")
                formatted.append(f"   Entity Type: {entity_type_display} | Related memories: {related_count}")
            
            return "\n".join(formatted)
            
        except Exception as e:
            logger.error(f"Entity search failed for '{entity_name}': {e}")
            return f"Entity search temporarily unavailable: {str(e)[:100]}"