import streamlit as st
from transformers import pipeline
import google.generativeai as genai
from pytube import Search
import speech_recognition as sr
import tempfile
from pydub import AudioSegment
import numpy as np
from streamlit_webrtc import webrtc_streamer, WebRtcMode, AudioProcessorBase

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

# 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 speech_to_text(audio_bytes):
    # Initialize recognizer
    r = sr.Recognizer()
    
    # Create a temporary file to store the recorded audio
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as fp:
        # Save audio bytes to temporary file
        fp.write(audio_bytes)
        temp_filename = fp.name
        
        # Read the audio file
        with sr.AudioFile(temp_filename) as source:
            # Adjust for ambient noise and record
            r.adjust_for_ambient_noise(source)
            audio = r.record(source)
            
            try:
                # Use Google Speech Recognition
                text = r.recognize_google(audio)
                return text
            except sr.UnknownValueError:
                st.error("Could not understand the audio")
                return None
            except sr.RequestError:
                st.error("Could not request results from speech recognition service")
                return None

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 option to choose between text and speech input
input_method = st.sidebar.radio("Choose input method:", ["Text", "Speech"])

if input_method == "Text":
    # Text input
    user_mood = st.sidebar.selectbox("Select your mood:", mood_options)
else:
    # Speech input
    st.write("📢 Tell me about your day...")
    
    # Use streamlit-webrtc for audio recording
    webrtc_ctx = webrtc_streamer(
        key="speech-to-text",
        mode=WebRtcMode.SENDONLY,
        audio_receiver_size=1024,
        media_stream_constraints={"audio": True, "video": False},
    )
    
    if webrtc_ctx.audio_receiver:
        audio_frames = webrtc_ctx.audio_receiver.get_frames(timeout=5)
        if audio_frames:
            audio_bytes = b"".join([frame.to_ndarray().tobytes() for frame in audio_frames])
            spoken_text = speech_to_text(audio_bytes)
            
            if spoken_text:
                st.write(f"You said: {spoken_text}")
                user_mood = detect_mood(spoken_text)
            else:
                user_mood = None
        else:
            st.warning("No audio frames received. Please try again.")
            user_mood = None
    else:
        user_mood = None

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

# Main content
if (user_mood or spoken_text) and gemini_api_key:
    mood = detect_mood(user_mood if user_mood else spoken_text)
    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)