import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import requests
import torch

# =======================
# RTL & UI Styling
# =======================
st.markdown(
    """
    <style>
    html, body, [class*="css"] {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }

    .stApp header, h1, h2, h3, h4, h5, h6 {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }

    .stMarkdown, .stMarkdown p {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }

    .stTextArea, .stTextInput, .stButton {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }

    [data-testid="stAlert"], [data-testid="stSpinnerText"] {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }

    section[data-testid="stSidebar"] {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Arial', sans-serif !important;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# =======================
# Classification Model
# =======================
MODEL_NAME = "imaneumabderahmane/Arabertv02-classifier-FA"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
LABELS = ["LABEL_0", "LABEL_1"]

def classify_question(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    pred = torch.argmax(outputs.logits, dim=1).item()
    return LABELS[pred]

# =======================
# LLM first-aid check
# =======================
def llm_first_aid_check(query, api_key):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    prompt = f"""
هل السؤال التالي يندرج ضمن الإسعافات الأولية؟ أجب فقط بـ "نعم" أو "لا":
{query}
"""
    data = {
        "model": "mistralai/mistral-7b-instruct",
        "messages": [
            {"role": "system", "content": "أنت مساعد ذكي عربي يساعد في تصنيف الأسئلة الطبية."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0,
        "max_tokens": 10
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        answer = response.json()["choices"][0]["message"]["content"].strip()
        return "نعم" in answer
    else:
        return False

# =======================
# Mistral Response 
# =======================
def get_mistral_response(prompt, api_key):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    few_shot_examples = """
أنت مساعد طبي محترف متخصص في الإسعافات الأولية. أجب دائمًا بنصائح إسعاف أولية واضحة ودقيقة وآمنة، باللغة العربية فقط.

سؤال 1: أصابتني ضربة على عظمة حاجبي الأيسر وحدوث انتفاخ بسيط. وضعت كمادات باردة. ماذا ينبغي أن أفعل أيضًا؟
جواب 1: استمر بوضع كمادات باردة لفترات قصيرة (10–15 دقيقة كل ساعة) لتخفيف الورم. ارفع رأسك عند الاستلقاء لتقليل التورم، وتجنّب الضغط على المنطقة. راقب أي تغيّر في الرؤية، فقدان للوعي، أو تقيؤ — إذا ظهرت أي من هذه الأعراض فراجع الطوارئ فورًا.

سؤال 2: كيف أتعامل مع شخص يتعرض لنوبة صرع؟
جواب 2: أبعد حوله الأشياء الحادة وحافظ على الهدوء. ضع شيئًا ناعمًا تحت رأسه، لا تُقيد حركاته، ولا تضع شيئًا في فمه. بعد انتهاء النوبة، ضع الشخص على جنبه للحفاظ على مجرى الهواء وراقبه حتى يستعيد وعيه. إذا استمرت النوبة أكثر من 5 دقائق أو تكررت، اطلب الإسعاف فورًا.
"""

    user_instruction = f"""
سؤال: {prompt}

اكتب الجواب بشكل مختصر ومباشر دون تكرار أو إعادة صياغة.
الجواب:
"""

    data = {
        "model": "mistralai/mistral-7b-instruct",
        "messages": [
            {"role": "system", "content": few_shot_examples},
            {"role": "user", "content": user_instruction}
        ],
        "temperature": 0.3,
        "max_tokens": 512
    }

    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"].strip()
    else:
        return f"⚠️ خطأ من API: {response.text}"

# =======================
# Streamlit UI
# =======================
st.title("المساعد الخاص بالإسعافات الأولية")

api_key = st.text_input("من فضلك أدخل المفتاح الخاص بك:", type="password")
user_input = st.text_area("اكتب سؤالك باللغة العربية:", height=100)

if st.button("اسأل"):
    if not api_key:
        st.warning("من فضلك أدخل المفتاح الخاص بك")
    elif not user_input.strip():
        st.warning("من فضلك اكتب سؤالك")
    else:
        with st.spinner("جارٍ تحليل سؤالك..."):
            category = classify_question(user_input)

        if category == "LABEL_1":
            with st.spinner("جارٍ التحقق من صحة السؤال..."):
                is_first_aid = llm_first_aid_check(user_input, api_key)

            if is_first_aid:
                st.success("تم التأكد من أن السؤال يتعلق بالإسعافات الأولية ✅")
                with st.spinner("جارٍ التواصل مع المساعد الذكي..."):
                    answer = get_mistral_response(user_input, api_key)
                    st.write("### 💬 الجواب:")
                    st.write(answer)
            else:
                st.warning("يبدو أن السؤال لا يتعلق بالإسعافات الأولية ❌")
        else:
            st.error("عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية❌")