from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, OpenAIServerModel, tool
import datetime
from superhero_party_theme_tool import SuperheroPartyThemeTool
from model import model

@tool
def suggest_menu(occasion: str) -> str:
    """
    Suggests a menu based on the occasion.
    Args:
        occasion: The type of occasion for the party. Allowed values are:
                  - "casual": Menu for casual party.
                  - "formal": Menu for formal party.
                  - "superhero": Menu for superhero party.
                  - "custom": Custom menu.
    """
    if occasion == "casual":
        return "Pizza, snacks, and drinks."
    elif occasion == "formal":
        return "3-course dinner with wine and dessert."
    elif occasion == "superhero":
        return "Buffet with high-energy and healthy food."
    else:
        return "Custom menu for the butler."


def main_v1() -> None:
    # Example 1: search the web for a party playlist
    music_agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model)
    music_agent.run("Search for the best music recommendations for a party at the Wayne's mansion.")

    # Example 2: use the custom menu tool
    menu_agent = CodeAgent(tools=[suggest_menu], model=model)
    menu_agent.run("Prepare a formal menu for the party.")

def main_v2() -> None:
    agent = CodeAgent(tools=[], model=model, additional_authorized_imports=['datetime'])

    agent.run(
        """
        Alfred needs to prepare for the party. Here are the tasks:
        1. Prepare the drinks - 30 minutes
        2. Decorate the mansion - 60 minutes
        3. Set up the menu - 45 minutes
        4. Prepare the music and playlist - 45 minutes

        If we start right now, at what time will the party be ready?
        """
    )

def main() -> None:
    agent = CodeAgent(tools=[SuperheroPartyThemeTool()], model=model)

    agent.run(
        """
        whats the theme for classic heroes party
        """
    )


if __name__ == "__main__":
    main()