import numpy as np
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# التحقق من توفر GPU
device = "cuda" if torch.cuda.is_available() else "cpu"

@st.cache_resource
def load_model():
    """
    تحميل النموذج والمُرمِّز مع التخزين المؤقت
    """
    model = AutoModelForCausalLM.from_pretrained(
        "joermd/speedy-llama2",
        torch_dtype=torch.float16,
        device_map=device
    )
    tokenizer = AutoTokenizer.from_pretrained("joermd/speedy-llama2")
    return model, tokenizer

# الصور العشوائية للكلاب عند حدوث خطأ
random_dog = [
    "0f476473-2d8b-415e-b944-483768418a95.jpg",
    "1bd75c81-f1d7-4e55-9310-a27595fa8762.jpg",
    "526590d2-8817-4ff0-8c62-fdcba5306d02.jpg",
    "1326984c-39b0-492c-a773-f120d747a7e2.jpg",
    "42a98d03-5ed7-4b3b-af89-7c4376cb14c3.jpg"
]

def reset_conversation():
    '''
    إعادة تعيين المحادثة
    '''
    st.session_state.conversation = []
    st.session_state.messages = []
    return None

def generate_response(prompt, temperature, max_length):
    """
    توليد استجابة من النموذج
    """
    try:
        inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
        
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=max_length,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    except Exception as e:
        return str(e)

# تحميل النموذج والمُرمِّز
try:
    with st.spinner('جاري تحميل النموذج... قد يستغرق هذا بضع دقائق...'):
        model, tokenizer = load_model()
except Exception as e:
    st.error(f"حدث خطأ أثناء تحميل النموذج: {str(e)}")
    st.stop()

# إعداد واجهة Streamlit
st.subheader('Mistral Chat')

# إضافة أزرار التحكم في الشريط الجانبي
temp_values = st.sidebar.slider('اختر قيمة درجة الحرارة', 0.0, 1.0, 0.5)
max_token_value = st.sidebar.slider('اختر الحد الأقصى للرموز', 100, 2000, 500)
st.sidebar.button('إعادة تعيين المحادثة', on_click=reset_conversation)

# تهيئة سجل المحادثة
if "messages" not in st.session_state:
    st.session_state.messages = []

# عرض رسائل المحادثة السابقة
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# معالجة إدخال المستخدم
if prompt := st.chat_input("اسألني سؤالاً"):
    # عرض رسالة المستخدم
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # عرض رد المساعد
    with st.chat_message("assistant"):
        try:
            # توليد الرد
            with st.spinner('جاري التفكير...'):
                response = generate_response(
                    prompt,
                    temperature=temp_values,
                    max_length=max_token_value
                )
            st.write(response)

        except Exception as e:
            response = "😵‍💫 يبدو أن هناك خطأ ما!\n حاول مرة أخرى لاحقاً.\n\n إليك صورة عشوائية لكلب 🐶:"
            st.write(response)
            random_dog_pick = 'https://random.dog/' + random_dog[np.random.randint(len(random_dog))]
            st.image(random_dog_pick)
            st.write("رسالة الخطأ:")
            st.write(e)

        st.session_state.messages.append({"role": "assistant", "content": response})