import streamlit as st
from transformers import pipeline
import google.generativeai as genai
from pytube import Search

# Load sentiment analysis model using PyTorch backend
mood_classifier = pipeline("sentiment-analysis", framework="pt")  # Use PyTorch

# Functions
def detect_mood(text):
    result = mood_classifier(text)[0]
    if result['label'] == 'POSITIVE':
        return "joyful"
    elif result['label'] == 'NEGATIVE':
        return "sad"
    else:
        return "neutral"

def get_song_recommendations(mood, api_key):
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-pro')

        # System prompt to guide the AI
        system_prompt = """
        You are a music recommendation assistant specialized in Indian songs. Your task is to suggest popular Indian songs based on the user's mood.
        - If the mood is "happy," suggest upbeat and joyful Bollywood or Indian pop songs.
        - If the mood is "sad," suggest emotional and soulful Indian songs.
        - If the mood is "energetic," suggest high-energy dance or workout songs.
        - If the mood is "romantic," suggest romantic Bollywood or Indian love songs.
        - If the mood is "calm," suggest soothing Indian classical or instrumental music.
        Always suggest 3 songs and provide the song title and artist name in the format:
        1. Song Title - Artist Name
        2. Song Title - Artist Name
        3. Song Title - Artist Name
        """
        
        # User prompt
        user_prompt = f"Suggest 3 popular Indian {mood} songs."
        
        # Combine system and user prompts
        response = model.generate_content([system_prompt, user_prompt])
        return response.text
    except Exception as e:
        st.error(f"Error using Gemini API: {e}")
        return None

def search_youtube(query):
    search = Search(query)
    return search.results[0].watch_url

# Streamlit App
st.title("🎵 Mood-Based Indian Song Player")

# Sidebar for user input
st.sidebar.header("How are you feeling today?")
mood_options = ["happy", "sad", "energetic", "romantic", "calm"]

# Input for Gemini API key
gemini_api_key = st.sidebar.text_input("Enter your Gemini API Key:", type="password")

# Add a button for mood input
user_mood = st.sidebar.selectbox("Select your mood:", mood_options)

# Playlist
if 'playlist' not in st.session_state:
    st.session_state.playlist = []

# Main content
if user_mood and gemini_api_key:
    mood = detect_mood(user_mood)
    st.write(f"🎭 Detected Mood: **{mood}**")
    
    st.write("🎵 Recommended Songs:")
    recommendations = get_song_recommendations(mood, gemini_api_key)
    if recommendations:
        st.write(recommendations)
        
        song_names = recommendations.split("\n")
        for song in song_names:
            if song.strip():
                st.write(f"🔍 Searching for: **{song}**")
                video_url = search_youtube(song)
                st.video(video_url)
                
                # Add to playlist button
                if st.button(f"Add '{song}' to Playlist"):
                    st.session_state.playlist.append((song, video_url))
                    st.success(f"Added '{song}' to your playlist!")
    else:
        st.error("Failed to get song recommendations. Please check your API key.")
elif not gemini_api_key:
    st.warning("Please enter your Gemini API key to get started.")

# Display the playlist
st.sidebar.header("Your Playlist")
for idx, (song, url) in enumerate(st.session_state.playlist):
    st.sidebar.write(f"{idx + 1}. {song}")
    if st.sidebar.button(f"Play {song}", key=f"play_{idx}"):
        st.video(url)