"""Wrapper singleton pour Qwen2.5-VL-3B.

Qwen2.5-VL-3B surpasse GLM-OCR sur les fiches OGC : extrait `dp_libelle`,
`praticien_conseil` (manuscrit !), `codage_reco.dp` et les 4 GHM/GHS là où
GLM-OCR échouait systématiquement. Un poil plus rapide aussi (3s vs 4s/page).

Coût : ~7 Go VRAM en bfloat16 (GLM-OCR = 2.2 Go) → tient sur RTX 5070 12 Go.
"""
import time
from pathlib import Path
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info

MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"


class QwenVLOCR:
    """Charge Qwen2.5-VL-3B une fois, réutilise le modèle pour toutes les pages."""
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._init_model()
        return cls._instance

    def _init_model(self):
        t0 = time.time()
        # max_pixels limite le nombre de patches visuels pour éviter l'OOM
        # sur images 300 dpi (2481x3509). ~1.25M pixels = équilibre qualité/VRAM.
        self.processor = AutoProcessor.from_pretrained(
            MODEL_PATH,
            min_pixels=256 * 28 * 28,
            max_pixels=1280 * 28 * 28,
        )
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            MODEL_PATH,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self.model.eval()
        self.load_time = time.time() - t0
        self.vram_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0

    def run(self, image_path: str | Path, prompt: str, max_new_tokens: int = 2048) -> dict:
        """Exécute Qwen-VL sur une image avec un prompt, retourne {text, elapsed_s}."""
        image_path = str(image_path)
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image_path},
                {"type": "text", "text": prompt},
            ],
        }]
        t0 = time.time()
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = self.processor(
            text=[text], images=image_inputs, videos=video_inputs,
            padding=True, return_tensors="pt",
        ).to(self.model.device)

        with torch.no_grad():
            generated_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
        out_ids = generated_ids[:, inputs.input_ids.shape[1]:]
        output = self.processor.batch_decode(out_ids, skip_special_tokens=True)[0]
        return {"text": output.strip(), "elapsed_s": time.time() - t0}