import os
import yaml
from Gradio_UI import GradioUI
from litellm import completion
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    FinalAnswerTool,
    LiteLLMModel,
    VisitWebpageTool,
    tool,
)

os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")

@tool
def analyze_lyrics_tool(song_title: str, artist: str, lyrics: str) -> str:
    """
    Performs a deep analysis of the musical track, given its metadata.
    
    Args:
        song_title: title of the song or music trach.
        artist: The name of the artist.
        lyrics: The lyrics of the song.
    
    Returns:
        A summary of the song's meaning in English.
    """

    prompt = f'''You are an expert in songs and their meanings. 
    Summarize the meaning of {song_title} by {artist} and identify 
    key themes based on the lyrics: 
    {lyrics}. 

    Includs deep idea and vibes analysis with explainations 
    based on references to the exact lines
    '''

    response = completion(
        model="gemini/gemini-2.0-flash", 
        messages=[
            {"role": "user", "content": prompt}
            ])
    
    try:
        lyrics = response.choices[0].message.content.strip()
        return lyrics
    except (AttributeError, KeyError, IndexError):
        try:
            lyrics = response['choices'][0]['message']['content'].strip()
            return lyrics
        except (AttributeError, KeyError, IndexError):
            pass

final_answer = FinalAnswerTool()
model = LiteLLMModel(model_id="gemini/gemini-2.0-flash")
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Example usage within the agent
agent = CodeAgent(
    tools=[
      FinalAnswerTool(), 
      DuckDuckGoSearchTool(),
      VisitWebpageTool(),
      analyze_lyrics_tool
    ], 
    model=model,
    additional_authorized_imports=['numpy', 'bs4'],
    max_steps=22,
    verbosity_level=0,
    name="Song Lyrics Analyzer",
    description="Analyze the meaning of song lyrics and identify key themes based on the lyrics using web search tools and deep lyrics analysis tool.",
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()