import streamlit as st
from transformers import pipeline
import random
import re
from textblob import TextBlob

# 1) Load a pretrained emotion classifier
@st.cache(allow_output_mutation=True)
def load_emotion_pipeline():
    return pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

clf = load_emotion_pipeline()

# 2) Your response logic (from your script)
responses = {
    "sadness": [
        "It’s okay to feel down sometimes. I’m here to support you.",
        "I'm really sorry you're going through this. Want to talk more about it?",
        "You're not alone — I’m here for you."
    ],
    "anger": [
        "That must have been frustrating. Want to vent about it?",
        "It's okay to feel this way. I'm listening.",
        "Would it help to talk through it?"
    ],
    "love": [
        "That’s beautiful to hear! What made you feel that way?",
        "It’s amazing to experience moments like that.",
        "Sounds like something truly meaningful."
    ],
    "joy": [  # note: pipeline uses 'joy' instead of 'happiness'
        "That's awesome! What’s bringing you joy today?",
        "I love hearing good news. 😊",
        "Yay! Want to share more about it?"
    ],
    "surprise": [
        "Oh! That’s surprising. Tell me more!",
        "Wow, that’s unexpected!",
        "Sounds like something caught you off guard."
    ],
    "fear": [
        "I hear your concern. 😟",
        "It’s okay to feel anxious sometimes.",
        "I’m here with you."
    ],
    "disgust": [
        "That sounds upsetting. I'm here to listen.",
        "I understand that’s hard to hear.",
        "Let’s talk through what’s bothering you."
    ],
    "neutral": [
        "Got it. I’m here if you want to dive deeper.",
        "Thanks for sharing that. Tell me more if you’d like.",
        "I’m listening. How else can I support you?"
    ]
}

help_keywords   = ["suggest","help","calm","exercise","relax","any tips","can you"]
bye_inputs      = ["bye","goodbye","exit","quit"]
thank_you_inputs= ["thank","thanks","thank you"]
awaiting_tip    = False
relaxation_resources = {
    "exercise": "Try the 5-4-3-2-1 grounding method:\n- 5 things you see\n- 4 you can touch\n- 3 you hear\n- 2 you smell\n- 1 you taste",
    "video":      "Here’s a short calming video: https://youtu.be/O-6f5wQXSu8"
}

# 3) Streamlit UI
st.title("🌿 EmotiBot")
user_input = st.text_input("You:")

if st.button("Send") and user_input:
    # Optional: spelling correction
    user = str(TextBlob(user_input).correct())

    # Exit
    if any(b in user.lower() for b in bye_inputs):
        st.write("EmotiBot 🌿: Take care! I’m here whenever you want to talk.")
    # Thank you
    elif any(t in user.lower() for t in thank_you_inputs):
        st.write("EmotiBot 🌿: You're most welcome! 💙")
    # Help keywords → ask tip
    elif any(k in user.lower() for k in help_keywords):
        awaiting_tip = True
        st.write("EmotiBot 🌿: Would you prefer a short calming video or a quick breathing exercise?")
    # Tip follow-up
    elif awaiting_tip:
        choice = "video" if "video" in user.lower() else "exercise"
        response = relaxation_resources.get(choice, relaxation_resources["exercise"])
        st.write(f"EmotiBot 🌿: {response}")
        awaiting_tip = False
    else:
        # 4) Emotion inference
        pred = clf(user)[0]["label"].lower()
        # Map pipeline labels to your keys
        pred = "happiness" if pred == "joy" else pred
        st.write(f"**Detected emotion:** {pred}")
        reply = random.choice(responses.get(pred, responses["neutral"]))
        st.write(f"EmotiBot 🌿: {reply}")