import os
import gradio as gr
from transformers import pipeline

# ============================================================
# 1. LOAD ARABERT CLASSIFIER
# ============================================================

print("Loading AraBERT classifier...")
CLF_MODEL = "imaneumabderahmane/Arabertv02-classifier-FA"
classifier = pipeline("text-classification", model=CLF_MODEL)
print("Classifier loaded successfully.")

# ============================================================
# 2. LOAD APOLLO GENERATOR
# ============================================================

print("Loading Apollo model...")
GEN_MODEL = "FreedomIntelligence/Apollo2-2B"

generator = pipeline(
    "text-generation",
    model=GEN_MODEL,
    torch_dtype="auto",
    device_map="auto"
)
print("Apollo loaded successfully.")

# ============================================================
# 3. GENERATION FUNCTION
# ============================================================

def generate_with_acegpt(prompt: str) -> str:
    """Generate a response in Arabic using AceGPT locally."""
    try:
        system_prompt = (
            "أنت مساعد طبي مختص في الإسعافات الأولية. "
            "قدّم إجابات دقيقة قصيرة و واضحة باللغة العربية الفصحى.\n\n"
        )
        input_text = system_prompt + f"المستخدم: {prompt}\nالمساعد:"
        result = generator(
            input_text,
            max_new_tokens=512,
            temperature=0.3,
            do_sample=True,
            top_p=0.9
        )
        return result[0]["generated_text"].split("المساعد:")[-1].strip()
    except Exception as e:
        print("Apollo generation error:", e)
        return "حدث خطأ أثناء توليد الإجابة من نموذج AceGPT."

# ============================================================
# 4. CHATBOT LOGIC
# ============================================================

def chatbot_fn(message: str, history: list):
    """Main function: classify → route → generate."""
    try:
        pred = classifier(message)[0]
        label = pred["label"]

        if label == "LABEL_1":
            response = generate_with_acegpt(message)
        else:
            response = "عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية."
    except Exception as e:
        print("Error in chatbot_fn:", e)
        response = "حدث خطأ أثناء معالجة الطلب."

    if history is None:
        history = []
    
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})
    
    return history, ""

# ============================================================
# 5. GRADIO INTERFACE
# ============================================================

with gr.Blocks(title="المساعد الذكي في الإسعافات الأولية") as demo:
    gr.Markdown(
        """
        # المساعد في الإسعافات الأولية  
        اكتب سؤالك بالعربية، وسيرد المساعد.  
        """
    )

    chatbot_ui = gr.Chatbot(
        label="المحادثة", 
        type="messages", 
        height=500,
        show_copy_button=True
    )
    
    with gr.Row():
        user_input = gr.Textbox(
            placeholder="اكتب سؤالك هنا...",
            label="سؤالك",
            lines=2,
            scale=8,
        )
        send_btn = gr.Button("إرسال", scale=1)
        clear_btn = gr.Button("مسح", scale=1)

    chat_state = gr.State([])

    send_btn.click(
        chatbot_fn,
        inputs=[user_input, chat_state],
        outputs=[chatbot_ui, user_input]
    )
    
    user_input.submit(
        chatbot_fn,
        inputs=[user_input, chat_state],
        outputs=[chatbot_ui, user_input]
    )
    
    clear_btn.click(
        lambda: ([], []),
        outputs=[chatbot_ui, chat_state]
    )

# ============================================================
# 6. LAUNCH
# ============================================================

if __name__ == "__main__":
    demo.launch()