import os
import gradio as gr
from transformers import pipeline
from openai import OpenAI

# ============================================================
# 1. ENVIRONMENT SETUP
# ============================================================

OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY")

if not OPENROUTER_KEY:
    raise ValueError("Missing OPENROUTER_API_KEY. Please add it to your Space secrets.")

client = OpenAI(
    api_key=OPENROUTER_KEY,
    base_url="https://openrouter.ai/api/v1"
)

# ============================================================
# 2. LOAD ARABERT CLASSIFIER
# ============================================================

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

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

def generate_with_deepseek(prompt: str) -> str:
    """Generate a response in Arabic using DeepSeek via OpenRouter."""
    try:
        completion = client.chat.completions.create(
            model="deepseek/deepseek-chat",
            messages=[
                {
                    "role": "system",
                    "content": (
                        "أنت مساعد طبي مختص في الإسعافات الأولية. "
                        "قدّم إجابات دقيقة وواضحة باللغة العربية الفصحى."
                    ),
                },
                {"role": "user", "content": prompt},
            ],
            max_tokens=512,
            temperature=0.3,
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        print("DeepSeek API error:", e)
        return "حدث خطأ أثناء الاتصال بنموذج DeepSeek عبر OpenRouter."

# ============================================================
# 4. CHATBOT LOGIC (FIXED FOR MESSAGES FORMAT)
# ============================================================

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_deepseek(message)
        else:
            response = "عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية."
    except Exception as e:
        print("Error in chatbot_fn:", e)
        response = "حدث خطأ أثناء معالجة الطلب."

    # For type="messages", append messages in the correct format
    if history is None:
        history = []
    
    # Append user message and assistant response in the correct format
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})
    
    return history, ""

# ============================================================
# 5. GRADIO INTERFACE (FIXED)
# ============================================================

with gr.Blocks(title="Chatbot الإسعافات الأولية") as demo:
    gr.Markdown(
        """
        # 🩺 Chatbot الإسعافات الأولية  
        اكتب سؤالك بالعربية، وسيرد المساعد بناءً على نموذج **DeepSeek** عبر OpenRouter.  
        """
    )

    chatbot_ui = gr.Chatbot(
        label="المحادثة", 
        type="messages", 
        height=500,
        show_copy_button=True  # Optional: adds copy button to messages
    )
    
    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)

    # State to maintain conversation history
    chat_state = gr.State([])

    # Actions
    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: ([], []),  # Clear both chatbot display and state
        outputs=[chatbot_ui, chat_state]
    )

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

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