import gradio as gr
import random
import os

from huggingface_hub import login
from smolagents import GradioUI, CodeAgent#, HfApiModel
# Try to import the HF Inference model class across smolagents versions
try:
    from smolagents.models import HfApiModel as HFInferenceModel
except Exception:
    try:
        from smolagents.models import InferenceAPIModel as HFInferenceModel
    except Exception:
        try:
            from smolagents.models import InferenceClientModel as HFInferenceModel
        except Exception:
            HFInferenceModel = None

# Import our custom tools from their modules
from tools import WeatherInfoTool, HubStatsTool#, DuckDuckGoSearchTool
from websearch import CustomDuckDuckGoSearchTool  as DuckDuckGoSearchTool
from retriever import load_guest_dataset

# Initialize the Hugging Face model
hf_token = os.getenv("HF_TOKEN")
login(hf_token)
hf_model = "Qwen/Qwen2.5-Coder-32B-Instruct" # HfApiModel default model

model = HFInferenceModel(token=hf_token)

# Initialize the web search tool
search_tool = DuckDuckGoSearchTool()

# Initialize the weather tool
weather_info_tool = WeatherInfoTool()

# Initialize the Hub stats tool
hub_stats_tool = HubStatsTool()

# Load the guest dataset and initialize the guest info tool
guest_info_tool = load_guest_dataset()

# Create Alfred with all the tools
alfred = CodeAgent(
    # tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
    tools=[guest_info_tool],
    model=model,
    add_base_tools=True,  # Add any additional base tools
    planning_interval=3   # Enable planning every 3 steps
)

if __name__ == "__main__":
    GradioUI(alfred).launch()