"""
RAG Tool
Search the knowledge base for relevant documents.
"""
from typing import Dict, Any, Optional
from loguru import logger

from tools.base import BaseTool, ToolResult


class RAGTool(BaseTool):
    """Search the RAG knowledge base."""
    
    def __init__(self, rag_store, config: Optional[Dict] = None):
        self.rag_store = rag_store
        super().__init__(config)
    
    @property
    def name(self) -> str:
        return "search_knowledge_base"
    
    @property
    def description(self) -> str:
        return "Search uploaded documents for relevant information. Use this for information from uploaded files, documents, or custom knowledge."
    
    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query"
                },
                "top_k": {
                    "type": "integer",
                    "description": "Number of results to return (default: 5)",
                    "default": 5
                }
            },
            "required": ["query"]
        }
    
    async def execute(self, query: str, top_k: int = 5, **kwargs) -> ToolResult:
        """Execute RAG search."""
        self._log_execution({"query": query, "top_k": top_k})
        
        try:
            results = await self.rag_store.search(query, top_k=top_k)
            
            if not results:
                return ToolResult(
                    success=True,
                    data="No relevant documents found in the knowledge base."
                )
            
            # Format results
            formatted_results = []
            for i, result in enumerate(results, 1):
                formatted_results.append(
                    f"{i}. From '{result.get('document_name', 'Unknown')}':\n"
                    f"   {result.get('content', '')}\n"
                    f"   Relevance: {result.get('score', 0):.2f}"
                )
            
            output = f"Knowledge base results for '{query}':\n\n" + "\n\n".join(formatted_results)
            
            self._log_success(output[:100])
            return ToolResult(success=True, data=output)
        
        except Exception as e:
            self._log_error(str(e))
            return ToolResult(success=False, error=str(e))