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 audio_recorder_streamlit import audio_recorder
import numpy as np
import wave
import io

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

def convert_audio_to_wav(audio_bytes):
    # Create a wave file in memory
    wav_buffer = io.BytesIO()
    
    with wave.open(wav_buffer, 'wb') as wav_file:
        wav_file.setnchannels(1)  # Mono
        wav_file.setsampwidth(2)  # 2 bytes per sample
        wav_file.setframerate(44100)  # Sample rate
        wav_file.writeframes(audio_bytes)
    
    return wav_buffer.getvalue()

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():
    # Initialize recognizer
    r = sr.Recognizer()
    
    # Add audio recording widget
    audio_bytes = audio_recorder(
        text="Click to record your mood",
        recording_color="#e8b62c",
        neutral_color="#6aa36f",
        pause_threshold=2.0  # Automatically stop after 2 seconds of silence
    )
    
    if audio_bytes:
        try:
            # Create a temporary WAV file
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
                # Convert and write audio bytes to WAV format
                wav_bytes = convert_audio_to_wav(audio_bytes)
                temp_audio.write(wav_bytes)
                temp_audio.flush()
                
                # Use the temporary file for speech recognition
                with sr.AudioFile(temp_audio.name) as source:
                    # Record the audio file
                    audio = r.record(source)
                    
                    try:
                        # Attempt speech recognition
                        text = r.recognize_google(audio)
                        st.success("Speech recognized successfully!")
                        return text
                    except sr.UnknownValueError:
                        st.error("Could not understand the audio. Please try speaking clearly and try again.")
                        return None
                    except sr.RequestError as e:
                        st.error(f"Could not request results from speech recognition service; {e}")
                        return None
        except Exception as e:
            st.error(f"Error processing audio: {e}")
            return None
    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"])

# Initialize user_mood as None
user_mood = None
user_text = None

if input_method == "Text":
    # Text input
    user_mood = st.sidebar.selectbox("Select your mood:", mood_options)
    user_text = user_mood
else:
    # Speech input
    st.write("📢 Click the button below and tell me about your day...")
    user_text = speech_to_text()
    
    if user_text:
        st.info(f"You said: '{user_text}'")
        user_mood = detect_mood(user_text)
        st.write(f"Based on what you said, I detect your mood as: **{user_mood}**")

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

# Main content
if user_text and gemini_api_key:
    detected_mood = detect_mood(user_text)
    st.write(f"🎭 Detected Mood: **{detected_mood}**")
    
    st.write("🎵 Recommended Songs:")
    recommendations = get_song_recommendations(detected_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)