import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# تنظیمات ادمین
ADMIN_USERNAME = "admin"
ADMIN_PASSWORD = "admin"

# CSS برای ظاهر برنامه
CUSTOM_CSS = """
<style>
body {
    font-family: 'Vazir', sans-serif;
    background: linear-gradient(135deg, #e8effa, #ffffff);
}
.chat-container {
    max-width: 400px;
    margin: 0 auto;
    border-radius: 15px;
    box-shadow: 0px 4px 12px rgba(0, 0, 0, 0.1);
}
.header {
    background: linear-gradient(135deg, #4a90e2, #50e3c2);
    color: white;
    padding: 10px;
    text-align: center;
    font-weight: bold;
    font-size: 18px;
}
.message {
    padding: 12px 15px;
    margin: 10px 0;
    border-radius: 15px;
    width: fit-content;
    max-width: 80%;
    font-size: 15px;
}
.message.user {
    align-self: flex-end;
    background: #4a90e2;
    color: white;
}
.message.bot {
    align-self: flex-start;
    background: #ff6b6b;
    color: white;
}
.footer {
    display: flex;
    padding: 10px;
    background: #f5f5f5;
    align-items: center;
}
.footer input[type="text"] {
    flex: 1;
    padding: 10px;
    border: none;
    border-radius: 20px;
    background: #ffffff;
    margin-right: 10px;
}
</style>
"""
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)

class ModernChatbot:
    def __init__(self):
        # تنظیم بارگذاری تنبل مدل
        if "model_loaded" not in st.session_state:
            st.session_state.model_loaded = False

    def load_model(self):
        """بارگذاری مدل و توکنایزر تنها در صورت نیاز"""
        if not st.session_state.model_loaded:
            model_name = "HuggingFaceH4/zephyr-7b-beta"
            st.session_state.tokenizer = AutoTokenizer.from_pretrained(model_name)
            st.session_state.model = AutoModelForCausalLM.from_pretrained(model_name)
            st.session_state.chat_pipeline = pipeline(
                "text-generation",
                model=st.session_state.model,
                tokenizer=st.session_state.tokenizer
            )
            st.session_state.model_loaded = True

    def generate_response(self, message):
        """تولید پاسخ با مدل بارگذاری شده"""
        self.load_model()
        response = st.session_state.chat_pipeline(
            message, max_length=50, num_return_sequences=1
        )[0]['generated_text']
        return response

    def login_admin(self, username, password):
        """ورود به عنوان مدیر"""
        return username == ADMIN_USERNAME and password == ADMIN_PASSWORD

# ذخیره تاریخچه مکالمات و وضعیت ورود ادمین در session_state
if "conversation_history" not in st.session_state:
    st.session_state.conversation_history = []

# ایجاد نمونه چت‌بات
chatbot = ModernChatbot()

# آیکون و پاپ‌آپ ورود ادمین
if st.sidebar.button("ورود ادمین"):
    username = st.sidebar.text_input("نام کاربری")
    password = st.sidebar.text_input("رمز عبور", type="password")
    if st.sidebar.button("تایید"):
        if chatbot.login_admin(username, password):
            st.session_state.admin_logged_in = True
            st.sidebar.success("ورود موفقیت‌آمیز بود")
        else:
            st.sidebar.error("نام کاربری یا رمز عبور اشتباه است")

# نمایش مکالمات
st.markdown("<div class='chat-container'><div class='header'>چت‌بات هوشمند</div>", unsafe_allow_html=True)
for sender, message in st.session_state.conversation_history:
    align_class = "user" if sender == "شما" else "bot"
    st.markdown(f"<div class='message {align_class}'>{sender}: {message}</div>", unsafe_allow_html=True)
st.markdown("</div>", unsafe_allow_html=True)

# ورودی پیام کاربر
user_message = st.text_input("پیام خود را اینجا بنویسید...", key="user_message")
if st.button("ارسال") and user_message:
    st.session_state.conversation_history.append(("شما", user_message))
    response = chatbot.generate_response(user_message)
    st.session_state.conversation_history.append(("ربات", response))
    st.experimental_rerun()

# پاک‌سازی مکالمات
if st.sidebar.button("پاک کردن تاریخچه مکالمات"):
    st.session_state.conversation_history = []
    st.experimental_rerun()