import streamlit as st
from streamlit_chat import message
from transformers import AutoModelForCausalLM, AutoTokenizer
import random

# Title and UI Customization
st.set_page_config(page_title="AI Chatbot", page_icon="🤖", layout="wide")
st.markdown(
    """
    <style>
    .main {background-color: #f5f7fa;}
    .stButton button {background-color: #6200EE; color: white; border-radius: 12px;}
    #chat-container {animation: fadeIn 2s ease-in;}
    @keyframes fadeIn {from {opacity: 0;} to {opacity: 1;}}
    </style>
    """,
    unsafe_allow_html=True,
)

# Load Hugging Face Model
def load_model():
    model_name = "gpt2"  # Replace with your preferred Hugging Face model
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return tokenizer, model

tokenizer, model = load_model()

# Store Chat History
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# User Interface
def chatbot_ui():
    st.markdown(
        """
        <div id="chat-container" style="text-align: center;">
            <h1>🚀 AI Chatbot</h1>
            <p>Ask me anything, and I'll do my best to help you!</p>
        </div>
        """,
        unsafe_allow_html=True,
    )
    user_input = st.text_input("Type your question:", "")


    if st.button("Send") and user_input:
        generate_response(user_input)

    # Display Chat History
    for i, chat in enumerate(st.session_state.chat_history):
        if chat['role'] == 'user':
            message(chat['content'], is_user=True, key=f'user_{i}')
        else:
            message(chat['content'], key=f'bot_{i}')

# Generate Response
def generate_response(user_input):
    inputs = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
    outputs = model.generate(inputs, max_length=200, num_return_sequences=1, do_sample=True, temperature=0.7)
    bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Update Chat History
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    st.session_state.chat_history.append({"role": "bot", "content": bot_reply})

chatbot_ui()