import os
from langchain_core.tools import tool
from pydantic import BaseModel, Field
from rag_with_memory import MemoryRAG
from loguru import logger

# Global variable to hold the engine
_rag_engine = None

def get_rag_engine():
    global _rag_engine
    if _rag_engine is None:
        # ABSOLUTE PATH for Docker 2026
        target_path = "/app/data/knowledge_base"
        
        # Verify folder exists and isn't empty
        if not os.path.exists(target_path) or not os.listdir(target_path):
            logger.error(f"KNOWLEDGE BASE MISSING at {target_path}")
            return None
            
        try:
            logger.info("Initializing MemoryRAG engine...")
            _rag_engine = MemoryRAG(docs_path=target_path)
            logger.success("MemoryRAG engine is now ONLINE.")
        except Exception as e:
            logger.error(f"MemoryRAG Initialization failed: {e}")
            return None
    return _rag_engine

class KnowledgeBaseInput(BaseModel):
    query: str = Field(description="The search query.")

@tool(args_schema=KnowledgeBaseInput, return_direct=True)
def knowledge_base_search(query: str) -> str:
    """Search the company knowledge base for answers."""
    engine = get_rag_engine()
    
    if engine is None:
        logger.error("Tool called but RAG engine remains None.")
        return "I'm sorry, I'm currently unable to access my internal documentation."

    try:
        # LOGIC FIX: Always return a string for the Agent to read
        result = engine.query(query, session_id="tool_session")
        return result.get("answer", "No information found.")
    except Exception as e:
        logger.error(f"Tool execution failed: {e}")
        return "I encountered an error while searching."