import streamlit as st
from transformers import pipeline

# Set up page title and icon
st.set_page_config(page_title="🎭 Roleplay Chatbot", page_icon="🤖")

# Load pre-trained chatbot model from Hugging Face
chatbot = pipeline("text-generation", model="nthakur/Mistral-7B-Instruct-v0.3-nomiracl-sft")

st.title("🎭 Roleplay AI Chatbot")
st.write("Chat with an AI character!")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state["messages"] = []

# Display chat history
for message in st.session_state["messages"]:
    avatar = "🧙‍♂️" if message["role"] == "assistant" else "🧑‍💻"
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

# User input field
user_input = st.chat_input("Type your message...")

if user_input:
    # Add user input to chat history
    st.session_state["messages"].append({"role": "user", "content": user_input})

    # Show typing animation
    with st.chat_message("assistant", avatar="🧙‍♂️"):
        with st.spinner("Typing..."):
            response = chatbot(user_input, max_length=100, do_sample=True, temperature=0.7)
            bot_reply = response[0]["generated_text"]
            st.markdown(bot_reply)

    # Save bot response to chat history
    st.session_state["messages"].append({"role": "assistant", "content": bot_reply})