import os
import gradio as gr
import torch
from PIL import Image
# Importation directe de la classe spécifique pour Qwen2.5-VL
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
import spaces  # Import crucial pour utiliser les GPU gratuits de Hugging Face

# --- CONFIGURATION ---
# REMPLACEZ 'votre_vrai_username' par votre véritable identifiant Hugging Face
MODEL_ID = "Ismailzeine/schema-bac-qwen-model"
HF_TOKEN = os.environ.get("HF_TOKEN")

print("⏳ Chargement du processeur et du modèle de vision Qwen2.5-VL...")
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
    token=HF_TOKEN
)
print("✅ Application prête !")


# --- FONCTION D'INFÉRENCE PROPULSÉE PAR ZEROGPU ---
@spaces.GPU  # Active le GPU gratuitement à la demande pour l'utilisateur
def run_inference(image_pil, question_text):
    if image_pil is None or not question_text.strip():
        return "Erreur : Veuillez fournir une image ET une question."

    # Préparation du prompt au format Qwen-VL
    messages = [
        {"role": "user", "content": [
            {"type": "image", "image": image_pil},
            {"type": "text", "text": question_text}
        ]}
    ]

    # Application du template de chat standard
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # Traitement des entrées pour le modèle (gère automatiquement les images et le texte)
    inputs = processor(text=[text], images=[image_pil], padding=True, return_tensors="pt").to("cuda")

    # Génération de la réponse
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False)

    # Découpage des tokens générés pour ne garder que la réponse de l'assistant
    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    response_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

    return response_text.strip()


# --- INTERFACE DESIGN GRADIO ---
logo_url = "https://imgur.com/GDKnbVp.png"

custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;600;700&display=swap');

.gradio-container {
    font-family: 'Poppins', sans-serif !important;
    background: linear-gradient(135deg, #f5f7fb 0%, #eef1f8 100%) !important;
    width: 100% !important;
    max-width: 100% !important; /* Modifié pour occuper tout l'espace disponible */
    margin: 0 auto !important;
    padding: 0 10px !important;
}

/* --- En-tête --- */
#header-block {
    background: linear-gradient(120deg, #4338CA 0%, #6366F1 55%, #4F46E5 100%);
    border-radius: 22px;
    padding: 34px 20px !important;
    margin-bottom: 26px;
    box-shadow: 0 12px 30px rgba(79, 70, 229, 0.28);
}
#header-block h1 {
    color: #ffffff !important;
    font-weight: 700 !important;
    letter-spacing: 0.3px;
}
#header-block p {
    color: #E0E7FF !important;
}

/* --- Cartes --- */
.card {
    background: #ffffff !important;
    border-radius: 18px !important;
    padding: 22px !important;
    box-shadow: 0 4px 22px rgba(15, 23, 42, 0.07) !important;
    border: 1px solid #EEF2FF !important;
}

.card h3, .card .prose h3 {
    margin-top: 0 !important;
    color: #312E81 !important;
    font-weight: 600 !important;
}

/* --- Bouton principal --- */
#submit-btn {
    background: linear-gradient(120deg, #4F46E5, #6366F1) !important;
    border: none !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    border-radius: 12px !important;
    box-shadow: 0 6px 16px rgba(79, 70, 229, 0.35) !important;
    transition: transform 0.15s ease, box-shadow 0.15s ease !important;
}
#submit-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 10px 24px rgba(79, 70, 229, 0.45) !important;
}

/* --- Zone de réponse --- */
#output-box textarea {
    background: #F8FAFC !important;
    border-radius: 12px !important;
    font-size: 1rem !important;
    border: 1px solid #E2E8F0 !important;
}

/* --- Exemples --- */
#examples-card .gr-samples-table, #examples-card table {
    border-radius: 12px !important;
    overflow: hidden !important;
}

footer {visibility: hidden}
"""

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue"),
    title="SCHÉMA BAC",
    css=custom_css,
    fill_width=True # Ajouté pour forcer l'interface à s'étaler horizontalement
) as demo:

    with gr.Column(elem_id="header-block"):
        gr.HTML(f'''
        <div style="text-align: center;">
           <img src="{logo_url}" alt="SCHÉMA BAC Logo" style="height: 70px; display: block; margin: 0 auto 12px auto;">
           
           <p style="font-size: 1.05rem; margin-top: 8px; max-width: 640px; margin-left:auto; margin-right:auto;">
               L'intelligence artificielle au service de la compréhension des schémas scientifiques.
           </p>
        </div>
        ''')

    with gr.Row(equal_height=True):
        with gr.Column(scale=1, min_width=320, elem_classes="card"): # min_width adapté pour maintenir la structure côte à côte
            gr.Markdown("### 1️⃣ Téléchargez votre schéma")
            image_input = gr.Image(type="pil", label="", height=280)

            gr.Markdown("### 2️⃣ Posez votre question")
            question_input = gr.Textbox(
                label="",
                placeholder="Ex : Quelle est la fonction de la structure numéro 3 ?",
                lines=3
            )

            submit_btn = gr.Button("🔍 Analyser le schéma", variant="primary", elem_id="submit-btn")

        with gr.Column(scale=1, min_width=320, elem_classes="card"): # min_width adapté pour maintenir la structure côte à côte
            gr.Markdown("### 🤖 Analyse de l'assistant")
            output_text = gr.Textbox(
                label="",
                interactive=False,
                lines=14,
                elem_id="output-box",
                placeholder="La réponse de l'IA apparaîtra ici..."
            )

    with gr.Group(elem_classes="card", elem_id="examples-card"):
        gr.Markdown("### 📚 Exemples à essayer — cliquez pour tester")
        gr.Examples(
            examples=[
                [
                    # 1. On remplace "examples" par "images"
                    # 2. Changez "votre_image_1.jpg" par le nom exact de votre premier fichier
                    os.path.join("images", "image.jpg"), 
                    "Quel est le rôle de la structure numéro 3 ?"
                ],
                [
                    # Faites de même pour la deuxième image
                    os.path.join("images", "image (3).jpg"),
                    "Que représente la structure numéro 3 sur cette coupe de rein ?"
                ],
                [
                    # Et pour la troisième image
                    os.path.join("images", "image (1).jpg"),
                    "Quelle est la fonction de la structure numéro 2 ?"
                ],
            ],
            inputs=[image_input, question_input],
            label=""
        )

    submit_btn.click(fn=run_inference, inputs=[image_input, question_input], outputs=output_text)

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