"""VLM fallback using Qwen2.5-VL for low-confidence product identification."""

import numpy as np
from PIL import Image

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

_model = None
_processor = None


def load_vlm(device: str = "cuda"):
    global _model, _processor
    if _model is not None:
        return _model, _processor

    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
    import torch

    print(f"[vlm] Loading Qwen2.5-VL ({QWEN_MODEL_ID}) on {device}...")
    _model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        QWEN_MODEL_ID,
        torch_dtype=torch.float16,
        device_map=device,
    )
    _processor = AutoProcessor.from_pretrained(QWEN_MODEL_ID)
    _model.eval()
    print("[vlm] Qwen2.5-VL loaded.")
    return _model, _processor


def identify(crop_bgr: np.ndarray, product_names: list[str], device: str = "cuda") -> str:
    """Use Qwen2.5-VL to identify a product crop from a list of known products.

    Returns the matched product name or 'unknown'.
    """
    import torch
    from qwen_vl_utils import process_vision_info

    model, processor = load_vlm(device)

    rgb = crop_bgr[:, :, ::-1]
    pil_img = Image.fromarray(rgb)

    product_list = "\n".join(f"- {name}" for name in product_names)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": pil_img},
                {"type": "text", "text": (
                    "This is a product on a retail shelf. Identify which product it is "
                    "from this list:\n"
                    f"{product_list}\n\n"
                    "If it matches one, reply with ONLY the exact product name from the list. "
                    "If it does not match any, reply with ONLY the word 'unknown'."
                )},
            ],
        }
    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to(device)

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=50)

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

    for name in product_names:
        if name.lower() in response.lower():
            return name

    return response


if __name__ == "__main__":
    import cv2
    import sys

    img_path = sys.argv[1] if len(sys.argv) > 1 else "outputs/crops/shelf_03_crop_000.jpg"
    crop = cv2.imread(img_path)

    product_names = [
        "Parle-G Gold Biscuits", "Britannia Good Day Butter Cookies",
        "Oreo Chocolate Creme Biscuits", "Sunfeast Dark Fantasy Choco Fills",
    ]

    result = identify(crop, product_names)
    print(f"VLM result: {result}")