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 typing import Optional
from duckduckgo_search import DDGS

# 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, model_name: str = "gemini-2.0-flash"):
        self.api_key = os.getenv("GOOGLE_API_KEY", "")  # Asume que está definida en Hugging Face Secrets
        self.model_name = model_name
        self.llm = setup_llm(self.api_key, self.model_name)
        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

        # Eliminamos las herramientas de smolagent_tools
        self.tools = [
            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:
        try:
            results = DDGS().text(query, max_results=3)
            return "\n".join([r["body"] for r in results])
        except Exception as e:
            return f"Error en 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)}"