"""
Medical Image Captioning — Qwen2.5-VL-7B + LoRA (correctness-focused checkpoint)

عشان ده يشتغل على HF Spaces مجانًا، بيستخدم ZeroGPU (GPU مشترك بيتاجَّر
بس وقت التنفيذ الفعلي). لازم الـ Space نفسه يتظبط بـ SDK: Gradio +
Hardware: ZeroGPU (مجاني، بس محتاج تفعيله من إعدادات الـ Space).
"""

import os
import spaces
import torch
import gradio as gr
from PIL import Image
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from peft import PeftModel

HF_TOKEN = os.environ.get("HF_TOKEN")  # اتحطت كـ Secret في إعدادات الـ Space

BASE_MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
ADAPTER_ID = "WafaaFraih/qwen25vl_correctness_checkpoint"

PROMPTS = {
    "Generic": "Describe this image.",
    "Clinical": (
        "You are a medical imaging expert. Describe this medical image "
        "with precise clinical terminology. Be concise and factual."
    ),
    "Chain-of-thought": (
        "You are a medical imaging expert. "
        "Step 1: Identify the imaging modality. "
        "Step 2: Identify the anatomical region. "
        "Step 3: Describe the key findings. "
        "Now write one concise clinical caption."
    ),
}

# ⚠️ ZeroGPU: مكتبة `spaces` بتعترض حتى عمليات تحميل الـ safetensors
# العادية لو حصلت برّه دالة معلَّمة بـ @spaces.GPU. عشان كده لازم
# التحميل كله (base model + adapter) يبقى جوه الدالة، مش هنا.
processor = None
model = None


def _ensure_loaded():
    """يحمّل الموديل والـ adapter مرة واحدة بس، أول ما حد يستخدم الـ demo."""
    global processor, model
    if model is not None:
        return
    print("▶ Loading base model + LoRA adapter (first request only)...")
    processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)
    base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.float16,
    )
    loaded = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
    loaded.eval()
    loaded.to("cuda")
    model = loaded
    print("✅ Model ready on GPU.")


@spaces.GPU(duration=120)
def caption_image(image: Image.Image, prompt_style: str):
    _ensure_loaded()

    if image is None:
        return "⚠️ Please upload a medical image first."

    image = image.convert("RGB")
    prompt_text = PROMPTS[prompt_style]

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt_text},
            ],
        }
    ]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda")

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
        )

    generated = processor.batch_decode(
        output_ids[:, inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )[0].strip()

    return generated


with gr.Blocks(title="Medical Image Captioning — Qwen2.5-VL (LoRA fine-tuned)") as demo:
    gr.Markdown(
        """
        # 🩺 Medical Image Captioning
        **Qwen2.5-VL-7B**, LoRA fine-tuned on ROCO + MultiCare with a
        correctness-focused training recipe (see the accompanying paper
        for the full evaluation, including known limitations).

        ⚠️ **Research demo only.** Captions are automatically generated
        and are **not validated for clinical or diagnostic use**. Do not
        use this tool to inform real medical decisions.
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload a medical image")
            prompt_choice = gr.Radio(
                choices=list(PROMPTS.keys()),
                value="Clinical",
                label="Prompt style",
            )
            run_button = gr.Button("Generate caption", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Generated caption", lines=6)

    run_button.click(fn=caption_image, inputs=[image_input, prompt_choice], outputs=output_text)

    gr.Markdown(
        """
        ---
        Model card and training details: see the linked paper and the
        [adapter repository](https://huggingface.co/WafaaFraih/qwen25vl_correctness_checkpoint).
        Base model: [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct).
        """
    )

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