import streamlit as st
from transformers import pipeline

# تحميل النموذج مرة واحدة فقط
@st.cache_resource
def load_generator():
    return pipeline("text-generation", model="aubmindlab/aragpt2-base", device=-1)  # استخدام CPU

# تحميل نموذج التوليد
generator = load_generator()

# عنوان التطبيق
st.title("مُحسِّن الجمل العربية والدردشة")

# --- قسم تحسين الجمل ---
st.subheader("تحسين جملة عربية")
user_input = st.text_input("أدخل جملة عربية لتحسينها:", "أنا ذهبت الحديقة")

if st.button("تحسين الجملة"):
    if user_input:
        # توليد الرد
        prompt = f"صحح هذه الجملة العربية وأجب فقط باللغة العربية: '{user_input}'"
        try:
            response = generator(prompt, max_new_tokens=50, temperature=0.7)[0]["generated_text"]
            # نحاول استخراج الجملة المصححة
            corrected_sentence = response.replace(prompt, "").strip()
            st.session_state.corrected_sentence = corrected_sentence
            st.success(f"✅ الجملة المحسّنة: {corrected_sentence}")
        except Exception as e:
            st.error(f"❌ حدث خطأ: {str(e)}")
    else:
        st.warning("⚠️ الرجاء إدخال جملة أولاً!")

# --- قسم الدردشة ---
st.subheader("الدردشة حول الجملة المحسّنة")

if "corrected_sentence" in st.session_state:
    chat_input = st.text_input("اطرح سؤالاً عن الجملة المحسّنة:", key="chat_input")
    if st.button("إرسال"):
        if chat_input:
            # توليد المحادثة
            prompt = (
                f"الجملة المصححة هي: '{st.session_state.corrected_sentence}'. "
                f"سؤال المستخدم: '{chat_input}'. "
                f"أجب فقط باللغة العربية الفصحى، ولا تستخدم أي كلمات إنجليزية."
            )
            try:
                response = generator(prompt, max_new_tokens=100, temperature=0.7)[0]["generated_text"]
                reply = response.replace(prompt, "").strip()
                st.write(f"🧑‍💬 **المستخدم:** {chat_input}")
                st.write(f"🤖 **المساعد:** {reply}")
            except Exception as e:
                st.error(f"❌ حدث خطأ أثناء المحادثة: {str(e)}")
        else:
            st.warning("⚠️ الرجاء كتابة سؤال!")
else:
    st.info("💬 يرجى تحسين جملة أولاً للبدء في الدردشة.")