from smolagents import CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, HfApiModel, Tool, tool, VisitWebpageTool
import os

@tool
def suggest_menu(occasion: str) -> str:
    """
    Предлагает меню на основе случая.
    Args:
        occasion: Тип случая для вечеринки (casual, formal, superhero).
    """
    if occasion == "casual":
        return "Пицца, закуски и напитки."
    elif occasion == "formal":
        return "3-ходовой ужин с вином и десертом."
    elif occasion == "superhero":
        return "Буфет с высокоэнергетической и здоровой едой."
    else:
        return "Пользовательское меню для дворецкого."

@tool
def catering_service_tool(query: str) -> str:
    """
    Этот инструмент возвращает самый рейтинговый сервис кейтеринга в Готэм-Сити.
    
    Args:
        query: Поисковый запрос для поиска сервисов кейтеринга.
    """
    # Примерный список сервисов кейтеринга и их рейтингов
    services = {
        "Gotham Catering Co.": 4.9,
        "Wayne Manor Catering": 4.8,
        "Gotham City Events": 4.7,
    }
    
    # Найти самый рейтинговый сервис кейтеринга (симуляция фильтрации поискового запроса)
    best_service = max(services, key=services.get)
    
    return best_service

class SuperheroPartyThemeTool(Tool):
    name = "superhero_party_theme_generator"
    description = """
    Этот инструмент предлагает творческие идеи вечеринок на тему супергероев на основе категории.
    понимает такие категории:
        "classic heroes"
        "villain masquerade"
        "futuristic Gotham"
    Он возвращает уникальную идею вечеринки."""
    
    inputs = {
        "category": {
            "type": "string",
            "description": "Тип вечеринки супергероев (например, 'классические герои', 'маскарад злодеев', 'футуристический Готэм').",
        }
    }
    
    output_type = "string"

    def forward(self, category: str):
        themes = {
            "classic heroes": "Гала Лиги Справедливости: Гости приходят одетыми в своих любимых героев DC с тематическими коктейлями, такими как 'Панч Криптонита'.",
            "villain masquerade": "Бал злодеев Готэма: Таинственный маскарад, где гости одеваются как классические злодеи Бэтмена.",
            "futuristic Gotham": "Ночь Нео-Готэма: Вечеринка в стиле киберпанка, вдохновленная Бэтменом из будущего, с неоновыми украшениями и футуристическими гаджетами."
        }
        
        return themes.get(category.lower(), "Идея тематической вечеринки не найдена. Попробуйте 'классические герои', 'маскарад злодеев' или 'футуристический Готэм'.")


# Альфред, дворецкий, готовит меню для вечеринки
agent = CodeAgent(
    tools=[
        DuckDuckGoSearchTool(), 
        VisitWebpageTool(),
        suggest_menu,
        catering_service_tool,
        SuperheroPartyThemeTool()
    ], 
    model=HfApiModel(token=os.getenv('hf_token')),
    max_steps=10,
    verbosity_level=2
)

agent.run("Дайте мне лучшее меню для вечеринки в особняке Уэйна. Идея вечеринки - тема 'маскарад злодеев'")