import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from gtts import gTTS
import random
from youtubesearchpython import Search

# Load DialoGPT model and tokenizer from Hugging Face
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Set up 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; }
    .css-ffhzg2 { font-size: 1.5em; font-weight: 500; color: #4C6D7D; }
    .stTextInput>div>div>input { background-color: #D8E3E2; }
    .stButton>button { background-color: #A9D0B6; color: white; border-radius: 5px; }
    .stButton>button:hover { background-color: #8FB79A; }
    .stTextInput>div>label { color: #4C6D7D; }
    </style>
""", unsafe_allow_html=True)

# Title and introduction to the bot
st.title("Grief and Loss Support Bot 🌿")
st.subheader("Your compassionate companion in tough times 💚")

# User input
user_input = st.text_input("Share what's on your mind...", placeholder="Type here...", max_chars=500)

# Store previous responses to check for repetition
if 'previous_responses' not in st.session_state:
    st.session_state.previous_responses = []

# Function to generate a more empathetic and focused response using DialoGPT
def generate_response(user_input):
    # Encode the input text and generate a response
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    bot_input_ids = new_user_input_ids
    chat_history_ids = model.generate(bot_input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id, temperature=0.7, top_k=50, repetition_penalty=1.2)
    
    # Decode the response to text
    chat_history_ids = chat_history_ids[:, bot_input_ids.shape[-1]:]  # remove the input from the response
    bot_output = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
    
    # Build a more empathetic and thoughtful response
    response = f"I’m really sorry you're feeling like this. It’s completely normal to feel overwhelmed when you're facing a heavy workload. It’s important to acknowledge how you feel and not keep it bottled up. Sometimes, stress and emotional exhaustion can build up, and it’s okay to let yourself feel those emotions."

    # Add coping strategies based on the situation
    if "workload" in user_input.lower():
        response += "\n\nWhen the workload feels too heavy, it can be helpful to break tasks down into smaller, more manageable steps. Focus on one thing at a time, and remember that it’s okay to take breaks when needed. Asking for support from colleagues or friends is also a good way to lighten the load."

    # Add general supportive message
    response += "\n\nYou're doing your best, and that’s all anyone can ask for. Please take care of yourself and know that it’s okay to take a step back when things feel too much. Your well-being is the most important thing."
    
    # Suggest a productive activity based on detected keywords
    if any(keyword in user_input.lower() for keyword in ["lonely", "lost", "sad", "overwhelmed", "academic", "exam"]):
        st.info("Here's a suggestion to help you cope:")

        # Providing a variety of activities based on user mood and needs
        activities = {
            "journaling": "Express your feelings in writing. Journaling is a great way to process emotions.",
            "yoga": "Yoga helps you relax and refocus. Try some deep breathing exercises or light stretching.",
            "painting": "Creative expression through painting or drawing can be soothing and help you release pent-up emotions.",
            "meditation": "Take a moment to calm your mind. Guided meditation can help reduce stress and anxiety.",
            "exercise": "Physical activity can lift your mood. Even a short walk in nature can make a big difference."
        }

        # Randomly select an activity category to suggest
        activity = random.choice(list(activities.keys()))
        st.write(f"How about {activity}? {activities[activity]}")

        # Search YouTube for videos related to the selected activity
        search = Search(activity)
        search_results = search.results[:3]  # limit results to 3 videos
        if search_results:
            for video in search_results:
                st.write(f"[{video.title}]({video.watch_url})")
        else:
            st.write("Sorry, I couldn't find any relevant videos at the moment.")

    # Crisis resources
    crisis_keywords = ["help", "suicide", "depressed", "emergency", "hurt", "lost"]
    if any(keyword in user_input.lower() for keyword in crisis_keywords):
        st.warning("It seems like you might be in distress. Please reach out to a crisis hotline or a trusted individual.")
        st.write("[Find emergency resources here](https://www.helpguide.org/find-help.htm)")

    return response

# Check if the user has typed something
if user_input:
    # Generate the empathetic response
    response = generate_response(user_input)
    
    # Store and show the new response
    st.session_state.previous_responses.append(response)
    st.text_area("Bot's Response:", response, height=250)

    # Text-to-speech output (optional)
    tts = gTTS(response, lang='en')
    audio_file = "response.mp3"
    tts.save(audio_file)
    st.audio(audio_file, format="audio/mp3")