import os
import requests
import streamlit as st
from transformers import pipeline
from PIL import Image
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Set up the Hugging Face API URL and your API key
API_URL = "https://api-inference.huggingface.co/models/trpakov/vit-face-expression"
headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}

# Function to query the Hugging Face model for facial expression
def query(filename):
    with open(filename, "rb") as f:
        data = f.read()
    response = requests.post(API_URL, headers=headers, data=data)
    
    if response.status_code == 200:
        return response.json()
    else:
        st.error("Error detecting facial expression: " + response.text)
        return None

# Function to generate a joke or uplifting text based on the mood
def generate_text_based_on_mood(emotion):
    try:
        # Predefined jokes for different emotions
        jokes = {
            "happy": [
                "Why did the scarecrow win an award? Because he was outstanding in his field!",
                "What do you call fake spaghetti? An impasta!"
            ],
            "sad": [
                "Why don’t scientists trust atoms? Because they make up everything!",
                "Why did the bicycle fall over? Because it was two-tired."
            ],
            "angry": [
                "Why did the tomato turn red? Because it saw the salad dressing!",
                "What did one angry man say to the other? Nothing, he just threw a rock!"
            ],
            "surprised": [
                "Did you hear about the claustrophobic astronaut? He just needed a little space!",
                "Why don’t skeletons fight each other? They don’t have the guts."
            ],
            "neutral": [
                "What do you call cheese that isn't yours? Nacho cheese!"
            ]
        }

        # Use predefined jokes if available
        if emotion in jokes:
            return jokes[emotion][0]  # Return the first joke for simplicity

        # If no predefined joke, use GPT-Neo
        generator = pipeline('text-generation', model='EleutherAI/gpt-neo-125M')
        prompt = f"Generate a light-hearted joke for someone who is feeling {emotion}."
        response = generator(prompt, max_length=50, num_return_sequences=1)
        return response[0]['generated_text']
        
    except Exception as e:
        st.error(f"Error generating text: {e}")
        return "Sorry, I couldn't come up with a joke at this moment."

# Function to convert text to speech using gTTS
def text_to_speech(text):
    from gtts import gTTS
    tts = gTTS(text, lang='en')
    audio_file = "output.mp3"
    tts.save(audio_file)  # Save the audio file
    return audio_file

# Streamlit UI
st.title("Facial Expression Mood Detector")
st.write("Upload an image of a face to detect mood and receive uplifting messages or jokes.")

# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load and display the image
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Save the uploaded file temporarily
    with open("uploaded_image.jpg", "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Detect facial expression
    expression_output = query("uploaded_image.jpg")
    if expression_output:
        # Assuming the response has a 'label' field with the detected emotion
        emotion = expression_output[0]['label']  # Adjust based on response structure
        st.write(f"Detected emotion: {emotion}")

        # Generate text based on detected emotion
        joke = generate_text_based_on_mood(emotion)
        st.write("Here's something to cheer you up:")
        st.write(joke)

        # Convert the generated joke to audio
        audio_file = text_to_speech(joke)

        # Provide an audio player in the Streamlit app
        st.audio(audio_file)  # Streamlit will handle playback