import os
import random
import zipfile
import subprocess
from transformers import pipeline
import gradio as gr

# Initialize Hugging Face translation pipeline
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")

# Initialize Hugging Face emotion classification pipeline
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

# Define emotion labels
emotion_labels = ["joy", "surprise", "sadness", "anger", "anxiety", "fear"]

# Base directory for storing music files (use a writable location)
base_dir = '/home/user/app/huggingface_models'

# Create directories for each emotion label
for emotion_label in emotion_labels:
    os.makedirs(f'{base_dir}/{emotion_label}', exist_ok=True)

# Download and unzip music files
subprocess.run([
    "wget", 
    "--no-check-certificate", 
    "https://github.com/AlanTFK/Cat/releases/download/NewAgeMusic/Music.zip", 
    "-O", 
    f'{base_dir}/Music.zip'
], check=True)

local_zip = f'{base_dir}/Music.zip'
zip_ref = zipfile.ZipFile(local_zip, 'r')
zip_ref.extractall(base_dir)
zip_ref.close()

# Function to classify emotions from text
def classify_emotions(text):
    translated_text = translator(text, max_length=512)[0]['translation_text']
    print(f"Translated text: {translated_text}")

    results = classifier(translated_text, return_all_scores=True)

    emotion_scores = {emotion: 0.0 for emotion in emotion_labels}
    for result in results:
        for score in result:
            emotion = score['label']
            if emotion in emotion_labels:
                emotion_scores[emotion] = score['score']

    return emotion_scores

# Function to get the main emotion from emotion scores
def get_main_emotion(emotion_scores):
    main_emotion = max(emotion_scores, key=emotion_scores.get)
    return main_emotion

# Function to select a random song from the emotion folder
def select_random_song(emotion_folder):
    files = os.listdir(emotion_folder)
    audio_files = [file for file in files if file.endswith('.mp3')]
    if not audio_files:
        raise FileNotFoundError(f"No audio files found in {emotion_folder}")

    random_song = random.choice(audio_files)
    return os.path.join(emotion_folder, random_song)

# Function to recommend music based on detected emotion
def emotion_music_recommendation(text):
    emotion_scores = classify_emotions(text)
    main_emotion = get_main_emotion(emotion_scores)
    emotion_folder = os.path.join(base_dir, main_emotion)

    try:
        selected_song = select_random_song(emotion_folder)
        return selected_song, f"Detected emotion: {main_emotion}"
    except FileNotFoundError as e:
        return str(e), "error"

# Gradio interface function
def play_music_interface(text):
    audio_file, emotion = emotion_music_recommendation(text)
    return audio_file, emotion

# Create Gradio interface
iface = gr.Interface(
    fn=play_music_interface,
    inputs="text",
    outputs=["audio", "text"],
    title="🎵 New Age AI 🎵",
    description="How are you feeling today?",
)

iface.launch()