import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from gtts import gTTS
from pytube import Search
import random
import os

# Load pretrained models
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
# Load GPT-2 model and tokenizer for story generation
gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
gpt2_model = AutoModelForCausalLM.from_pretrained("gpt2-medium")

emotion_classifier = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", return_all_scores=True)

# Function to generate a comforting story using GPT-2
def generate_story(theme):
    # A detailed prompt for generating a comforting story about the selected theme
    story_prompt = f"Write a comforting, detailed, and heartwarming story about {theme}. The story should include a character who faces a tough challenge, finds hope, and ultimately overcomes the situation with a positive resolution."
    
    # Generate story using GPT-2
    input_ids = gpt2_tokenizer.encode(story_prompt, return_tensors='pt')
    
    story_ids = gpt2_model.generate(
        input_ids,
        max_length=500,  # Generate longer stories
        temperature=0.7,  # Balanced creativity
        top_p=0.9,
        repetition_penalty=1.2,
        num_return_sequences=1
    )
    
    # Decode the generated text
    story = gpt2_tokenizer.decode(story_ids[0], skip_special_tokens=True)
    return story

def generate_response(user_input):
    # Refined prompt with a clear and empathetic tone
    response_prompt = f"You are a kind and empathetic support bot. A user is sharing their feelings: '{user_input}'. Respond with kindness and empathy, offering emotional validation. Keep the tone soft and comforting, and avoid any philosophical or unrelated explanations. Offer support and let them know that their feelings are valid."
    
    # Generate the response using the GPT-2 model
    input_ids = gpt2_tokenizer.encode(response_prompt, return_tensors='pt')
    response_ids = gpt2_model.generate(
        input_ids,
        max_length=300,
        temperature=0.85,
        top_k=50,
        repetition_penalty=1.2,
        num_return_sequences=1
    )
    
    # Decode the response and clean it up by removing the prompt
    response = gpt2_tokenizer.decode(response_ids[0], skip_special_tokens=True)
    
    # Strip out the prompt portion to get a clean, empathetic message
    cleaned_response = response.replace(f"You are a kind and empathetic support bot. A user is sharing their feelings: '{user_input}'. Respond with kindness and empathy, offering emotional validation. Keep the tone soft and comforting, and avoid any philosophical or unrelated explanations. Offer support and let them know that their feelings are valid.", "").strip()
    
    return cleaned_response



# Analyze user input for emotional tone
def get_emotion(user_input):
    emotions = emotion_classifier(user_input)
    emotions_sorted = sorted(emotions[0], key=lambda x: x['score'], reverse=True)
    return emotions_sorted[0]['label']

# Function to fetch YouTube videos
def fetch_youtube_videos(activity):
    search = Search(f"{activity} for mental health relaxation")
    search_results = search.results[:3]
    videos = []
    for video in search_results:
        video_url = f"https://www.youtube.com/watch?v={video.video_id}"
        videos.append((video.title, video_url))
    return videos

# Streamlit page configuration
st.set_page_config(page_title="Grief and Loss Support Bot 🌿", page_icon="🌿", layout="centered")
st.markdown("<style>.css-1d391kg { background-color: #F3F7F6; }</style>", unsafe_allow_html=True)

st.title("Grief and Loss Support Bot 🌿")
st.subheader("Your compassionate companion in tough times 💚")

# Sidebar for Meditation and Story Generation
with st.sidebar:
    st.header("🧘 Guided Meditation")
    if st.button("Play Meditation"):
        meditation_audio = "meditation.mp3"
        if not os.path.exists(meditation_audio):
            tts = gTTS("Take a deep breath. Relax and let go of any tension...", lang='en')
            tts.save(meditation_audio)
        st.audio(meditation_audio, format="audio/mp3")

# Generating a comforting story
st.sidebar.header("📖 Short Comforting Story")
story_theme = st.selectbox("Choose a theme for your story:", ["courage", "healing", "hope"])
if st.sidebar.button("Generate Story"):
    with st.spinner("Generating your story..."):
        story = generate_story(story_theme)
    st.text_area("Here's your story:", story, height=300)



# User input section
user_input = st.text_input("Share what's on your mind. I am here to listen...", placeholder="Type here...", max_chars=500, key="user_input_1")


# Initialize session state
if 'previous_responses' not in st.session_state:
    st.session_state.previous_responses = []
if 'badges' not in st.session_state:
    st.session_state.badges = []


# Initialize session state
if 'badges' not in st.session_state:
    st.session_state.badges = []

if user_input:
    with st.spinner("Analyzing your input..."):
        # Get the emotion of the user input
        emotion = get_emotion(user_input)
        
        # Generate an empathetic response
        response = generate_response(user_input)
        
        # Display the bot's response
        st.text_area("Bot's Response:", response, height=250)
    
    # Assign badges based on the detected emotion
    if emotion in ["joy", "optimism"]:
        badge = "🌟 Positivity Badge"
        if badge not in st.session_state.badges:
            st.session_state.badges.append(badge)
            st.success(f"Congratulations! You've earned a {badge}!")

    # Suggest activities based on emotion
    st.info("🎨 Try a New Activity")
    activities = ["exercise", "yoga", "journaling", "painting", "meditation", "swimming"]
    selected_activity = st.selectbox("Pick an activity:", activities)

    if st.button("Find Videos"):
        videos = fetch_youtube_videos(selected_activity)
        if videos:
            for title, url in videos:
                st.write(f"[{title}]({url})")
        else:
            st.write(f"No results found for '{selected_activity}'.")

# Crisis resources
if user_input and any(word in user_input.lower() for word in ["suicide", "help", "depressed"]):
    st.warning("Please reach out to a crisis hotline for immediate support.")
    st.write("[Find emergency resources here](https://www.helpguide.org/find-help.htm)")


# Generate audio response
if user_input:
    tts = gTTS(response, lang='en')
    audio_file = "response.mp3"
    tts.save(audio_file)
    st.audio(audio_file, format="audio/mp3")

# Display badgesz
if st.session_state.badges:
    st.sidebar.header("🏅 Achievements")
    for badge in st.session_state.badges:
        st.sidebar.write(badge)