import streamlit as st
from transformers import pipeline

# بارگذاری مدل از Hugging Face
@st.cache_resource
def load_model():
    generator = pipeline("text-generation", model="gpt2")  # یا هر مدل مناسب دیگر
    return generator

model = load_model()

def generate_response(prompt):
    response = model(prompt, max_length=100, num_return_sequences=1)  # تنظیم max_length برای پاسخ‌های طولانی‌تر
    return response[0]["generated_text"]

def main():
    st.set_page_config(page_title="دستیار هوش مصنوعی", layout="wide")
    st.title("دستیار هوش مصنوعی پیشرفته (Hugging Face)")

    # مقداردهی اولیه Session State
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # نمایش HTML
    with open("index.html", "r", encoding="utf-8") as file:
        html_content = file.read()

    # اضافه کردن کد Javascript برای ارتباط با Streamlit
    html_content = html_content.replace("</body>", """
        <script>
            // تابع ارسال پیام به Streamlit
            function sendMessage() {
                var userInput = document.getElementById("user-input").value;
                document.getElementById("user-input").value = ""; // پاک کردن ورودی
                // اضافه کردن پیام کاربر به لیست پیام‌ها
                var chatMessages = document.getElementById("chat-messages");
                var newMessage = document.createElement("div");
                newMessage.classList.add("message", "user-message");
                newMessage.innerText = userInput;
                chatMessages.appendChild(newMessage);
                // ارسال پیام به Streamlit
                Streamlit.setComponentValue("user-input", userInput);
            }

            // دریافت پاسخ از Streamlit و نمایش اون
            Streamlit.on('update', function() {
                var response = Streamlit.getComponentValue("response");
                if (response) {
                    // اضافه کردن پیام به لیست پیام‌ها
                    var chatMessages = document.getElementById("chat-messages");
                    var newMessage = document.createElement("div");
                    newMessage.classList.add("message", "assistant-message");
                    newMessage.innerText = response;
                    chatMessages.appendChild(newMessage);
                }
            });

            // ارسال پیام با فشردن Enter
            var userInput = document.getElementById("user-input");
            userInput.addEventListener("keyup", function(event) {
                if (event.keyCode === 13) {
                    event.preventDefault();
                    document.getElementById("send-button").click();
                }
            });
        </script>
    </body>
    """)

    st.components.v1.html(html_content, height=700, scrolling=True)

    # دریافت پیام از Session State
    user_input = st.components.v1.text_input("user-input", value="", key="user_input")
    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        response = generate_response(user_input)
        st.session_state.messages.append({"role": "assistant", "content": response})

    # ارسال پاسخ به Session State
    if st.session_state.messages:
        st.components.v1.text_area("response", value=st.session_state.messages[-1]["content"], key="response")

if __name__ == "__main__":
    main()