import streamlit as st
from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast, pipeline
import time
from streamlit_chat import message


import base64

# Function to set background
def add_bg_from_local(image_file):
    with open(image_file, "rb") as f:
        encoded = base64.b64encode(f.read()).decode()
    css = f"""
    <style>
    .stApp {{
        background-image: url("data:image/png;base64,{encoded}");
        background-size: cover;
        background-position: center;
        background-attachment: fixed;
    }}
    </style>
    """
    st.markdown(css, unsafe_allow_html=True)

# Call function with your image
add_bg_from_local("pexels-adrien-olichon-1257089-2387793.jpg")   # replace with your image file name

# ---------------------------
# Load fine-tuned model
# ---------------------------
model_path = "./models/sentiment"
model = DistilBertForSequenceClassification.from_pretrained(model_path)
tokenizer = DistilBertTokenizerFast.from_pretrained(model_path)
sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)

# ---------------------------
# Animated Background with Emojis
# ---------------------------
def sentiment_background(label):
    if label in ["POSITIVE", "VERY_POSITIVE"]:
        emoji = "😍"
        colors = ["#56ab2f", "#a8e063"]
    elif label == "NEUTRAL":
        emoji = "😐"
        colors = ["#bdc3c7", "#2c3e50"]
    else:
        emoji = "💔"
        colors = ["#e96443", "#904e95"]

    st.markdown(
        f"""
        <style>
        body {{
            background: linear-gradient(-45deg, {colors[0]}, {colors[1]});
            background-size: 400% 400%;
            animation: gradientBG 12s ease infinite;
        }}
        @keyframes gradientBG {{
            0% {{background-position: 0% 50%;}}
            50% {{background-position: 100% 50%;}}
            100% {{background-position: 0% 50%;}}
        }}
        .emoji {{
            position: fixed;
            font-size: 2rem;
            animation: float 6s infinite;
            opacity: 0.7;
        }}
        @keyframes float {{
            0% {{ transform: translateY(100vh) translateX(0); opacity: 0; }}
            50% {{ opacity: 1; }}
            100% {{ transform: translateY(-10vh) translateX(20px); opacity: 0; }}
        }}
        </style>
        <div class="emoji" style="left:10%;">{emoji}</div>
        <div class="emoji" style="left:30%;">{emoji}</div>
        <div class="emoji" style="left:50%;">{emoji}</div>
        <div class="emoji" style="left:70%;">{emoji}</div>
        <div class="emoji" style="left:90%;">{emoji}</div>
        """,
        unsafe_allow_html=True,
    )

# ---------------------------
# App Layout
# ---------------------------
st.set_page_config(page_title="Sentiment Chatbot", page_icon="🤖", layout="centered")

st.title("🤖 Sentiment-Aware Chatbot")
st.markdown("This chatbot understands **VERY NEGATIVE, NEGATIVE, NEUTRAL, POSITIVE, VERY POSITIVE** emotions.")

if "history" not in st.session_state:
    st.session_state.history = []

# ---------------------------
# Input box + Send button
# ---------------------------
col1, col2 = st.columns([4, 1])

with col1:
    user_input = st.text_input(
        "You:",
        placeholder="Type your message here...",
        key="input_box",
        label_visibility="collapsed"  # hides label so box aligns better
    )

with col2:
    send = st.button("Send")

# ---------------------------
# When user sends a message
# ---------------------------
if send and user_input:
    # Save user message
    st.session_state.history.append(("user", user_input))

    # Bot typing ani
    with st.spinner("🤖 Bot is typing..."):
        time.sleep(1.5)

        # Predict sentiment
        result = sentiment_pipeline(user_input)[0]
        label = result["label"]
        score = result["score"]

        # Update background with floating emojis
        sentiment_background(label)

        # Bot reply
        if label == "VERY_POSITIVE":
            bot_reply = f"🌟 Wow! That's fantastic! ({label}: {score:.2f})"
        elif label == "POSITIVE":
            bot_reply = f"😊 I'm glad to hear that! ({label}: {score:.2f})"
        elif label == "NEUTRAL":
            bot_reply = f"😐 Thanks for sharing your thoughts. ({label}: {score:.2f})"
        elif label == "NEGATIVE":
            bot_reply = f"😟 I'm sorry to hear that. ({label}: {score:.2f})"
        else:  # VERY_NEGATIVE
            bot_reply = f"💔 That sounds really bad. Stay strong! ({label}: {score:.2f})"

        st.session_state.history.append(("bot", bot_reply))

# ---------------------------
# Display conversation
# ---------------------------
for role, text in st.session_state.history:
    if role == "user":
        message(text, is_user=True, avatar_style="personas")
    else:
        message(text, avatar_style="bottts")

# ---------------------------
# Clear Chat Button
# ---------------------------
if st.button("Clear Chat"):
    st.session_state.history = []