import time
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentExecutor, Tool
from langchain.chat_models import ChatGoogleGenerativeAI
import genai
from gemini import GeminiAgent
from smolagent_tools import WikipediaSearchTool, WebSearchTool, SmolagentToolWrapper

# Función para configuración de llm
def setup_llm(api_key: str, model_name: str = "gemini-2.0-flash"):
    """Configuración del modelo de lenguaje con Gemini."""
    genai.configure(api_key=api_key)
    return ChatGoogleGenerativeAI(model_name=model_name)

# Herramientas adicionales
class EnhancedAgent:
    def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
        self.api_key = api_key
        self.model_name = model_name
        self.llm = setup_llm(api_key, model_name)
        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

        self.tools = [
            SmolagentToolWrapper(WikipediaSearchTool()),
            Tool(name="web_search", func=self._web_search, description="Realizar búsqueda web avanzada")
        ]

        self.agent = GeminiAgent(
            tools=self.tools,
            llm=self.llm,
            memory=self.memory,
        )

    def _web_search(self, query: str, domain: Optional[str] = None) -> str:
        """Realiza una búsqueda en la web usando DuckDuckGo con limitación de tasa y reintentos."""
        try:
            search_tool = WebSearchTool()
            results = search_tool.search(query, domain)
            return results
        except Exception as e:
            return f"Error al realizar la búsqueda: {str(e)}"

    def run(self, query: str) -> str:
        """Procesa las consultas del usuario con reintentos y manejo de errores."""
        max_retries = 3
        base_sleep = 1
        for attempt in range(max_retries):
            try:
                response = self.agent.run(query)
                return response
            except Exception as e:
                sleep_time = base_sleep * (attempt + 1)
                if attempt < max_retries - 1:
                    print(f"Intento {attempt + 1} fallido. Reintentando en {sleep_time} segundos...")
                    time.sleep(sleep_time)
                    continue
                return f"Error procesando la consulta: {str(e)}"