import gradio as gr
import random
import os
import getpass
from smolagents import GradioUI, CodeAgent, HfApiModel

# Import our custom tools from their modules
from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool, CrashInfoRetrieverTool
from retriever import load_guest_dataset, load_crash_data

#Set up in HuggingFace Space Secrets
#os.environ["HF_TOKEN"] = getpass.getpass("Enter HF token:")

# Ensure the HF_TOKEN environment variable is set
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise ValueError("HF_TOKEN environment variable is not set.")

# Initialize the Hugging Face model
model = HfApiModel() #assume HF will look for it

# 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()

crash_info_tool = load_crash_data()

# Create Alfred with all the tools
alfred = CodeAgent(
    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool, crash_info_tool], 
    model=model,
    add_base_tools=True,  # Add any additional base tools
    planning_interval=3   # Enable planning every 3 steps
    # memory=True  # Enable conversation memory.  NOTE: this wasn't supported for me..
)

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