import gradio as gr
import random
from smolagents import GradioUI, CodeAgent, LiteLLMModel
from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
from retriever import load_guest_dataset
import os
import logging
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import litellm

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class RetryLiteLLMModel(LiteLLMModel):
    @retry(
        stop=stop_after_attempt(3),  # Максимум 3 попытки
        wait=wait_exponential(multiplier=1, min=1, max=10),  # Задержка: 1с, 2с, 4с (макс. 10с)
        retry=retry_if_exception_type(litellm.InternalServerError),  # Повторять только для InternalServerError
        before=lambda _: logger.info("Retrying due to 503 error..."),
        after=lambda retry_state: logger.info(f"Retry attempt {retry_state.attempt_number} completed")
    )
    def generate(self, *args, **kwargs):
        """Переопределение метода generate с повторными попытками"""
        return super().generate(*args, **kwargs)

# Model wrapper
try:
    model = RetryLiteLLMModel(
        model_id="gemini/gemini-2.0-flash-lite",
        api_key=os.environ['GEMINI_API_TOKEN']
    )
    logger.info("Successfully initialized LiteLLM model")
except KeyError as e:
    logger.error("GEMINI_API_TOKEN not set in environment variables")
    raise
except Exception as e:
    logger.error(f"Failed to initialize model: {str(e)}")
    raise

# Tools initialization
search_tool = DuckDuckGoSearchTool()
weather_info_tool = WeatherInfoTool()
hub_stats_tool = HubStatsTool()
guest_info_tool = load_guest_dataset()

alfred = CodeAgent(
    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool],
    model=model,
    add_base_tools=True,
    planning_interval=3
)

if __name__ == "__main__":
    logger.info("Launching Gradio UI...")
    GradioUI(alfred).launch()