from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    HfApiModel,  # ✅ Use this instead
    load_tool,
    tool
)
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# 🎵 Mood-based Playlist Generator Tool
@tool
def mood_playlist_generator(mood: str) -> str:
    """Generates a music playlist based on the user's mood.
    Args:
        mood: The emotional state or context (e.g., 'focus', 'joy', 'relaxation').
    """
    playlists = {
        "focus": [
            "Deep Focus – Spotify: https://open.spotify.com/playlist/37i9dQZF1DX4sWSpwq3LiO",
            "Lo-Fi Beats – Spotify: https://open.spotify.com/playlist/37i9dQZF1DWWQRwui0ExPn"
        ],
        "joy": [
            "Have a Great Day! – Spotify: https://open.spotify.com/playlist/37i9dQZF1DX3rxVfibe1L0",
            "Feelin' Good – Spotify: https://open.spotify.com/playlist/37i9dQZF1DXdPec7aLTmlC"
        ],
        "relax": [
            "Peaceful Piano – Spotify: https://open.spotify.com/playlist/37i9dQZF1DX4sWSpwq3LiO",
            "Calm Vibes – Spotify: https://open.spotify.com/playlist/37i9dQZF1DWZqd5JICZI0u"
        ],
        "energy": [
            "Beast Mode – Spotify: https://open.spotify.com/playlist/37i9dQZF1DX76Wlfdnj7AP",
            "Power Workout – Spotify: https://open.spotify.com/playlist/37i9dQZF1DWUVpAXiEPK8P"
        ]
    }

    mood = mood.lower()
    if mood not in playlists:
        return "Sorry, I don’t have a playlist for that mood. Try: focus, joy, relax, or energy."

    return f"Here’s a playlist for when you’re feeling **{mood}**:\n" + "\n".join(playlists[mood])


# 🧠 LLM Setup
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# 🎨 Tool: Image Generation
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# 📄 Load system prompts
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# 🔧 Register tools
final_answer = FinalAnswerTool()
tools = [final_answer, image_generation_tool, mood_playlist_generator]

# 🤖 Create the Agent
agent = CodeAgent(
    model=model,
    tools=tools,
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="MoodTune",
    description="An agent that generates music playlists and images based on your mood.",
    prompt_templates=prompt_templates
)

# 🚀 Launch Gradio UI
GradioUI(agent).launch()