import chromadb
from chromadb.utils import embedding_functions
from smolagents import GradioUI, CodeAgent, HfApiModel

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

# Initialize the Hugging Face model
model = HfApiModel(max_tokens=8096)

# 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
chroma_client = chromadb.Client()
default_ef = embedding_functions.DefaultEmbeddingFunction()
collection = chroma_client.create_collection(name="guests-collection", embedding_function=default_ef)

guest_info_tool = load_guest_dataset(collection)

# Create Alfred with all the tools
search_agent = CodeAgent(
    name="search_agent",
    description="Browses the web to find information",
    tools=[search_tool],
    model=model,
    add_base_tools=False,
    planning_interval=3,
    verbosity_level=2,
    max_steps=5
)
hub_stats_agent = CodeAgent(
    name="hub_stats_agent",
    description="Browses the HF statistics",
    tools=[hub_stats_tool],
    model=model,
    add_base_tools=False,
    planning_interval=3,
    verbosity_level=2,
    max_steps=2
)
weather_agent = CodeAgent(
    name="weather_agent",
    description="Weather forecasting",
    tools=[weather_info_tool],
    model=model,
    add_base_tools=False,
    planning_interval=3,
    verbosity_level=2,
    max_steps=2
)

alfred_manager = CodeAgent(
    tools=[guest_info_tool],
    managed_agents=[search_agent, hub_stats_agent, weather_agent],
    model=model,
    add_base_tools=False,
    planning_interval=5,
    verbosity_level=2,
    max_steps=15
)

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