from PIL import Image, UnidentifiedImageError
from utils import *


def get_evaluation_data(ds):
    evaluation_data = []
    for i in range(len(ds)):
        try:
            img = ds[i]["image"]
            thumbnail_img = img.copy()
            thumbnail_img.thumbnail((256, 256))
            evaluation_data.append({
                "id": ds[i]["ex_id"],
                "image_thumbnail": image_to_base64(thumbnail_img),
                "image_full": image_to_base64(img),
                "image_full_url": "https://huggingface.co/", # Dummy
                "prompt": ds[i]["prompt"],
                "category": ds[i]["category"]
            })
        except (UnidentifiedImageError, OSError, ValueError): # To handle .heic images -> can be removed when dataset is fixed
            img = Image.new("RGB", (256, 256), color="white")
            evaluation_data.append({
                "id": i,
                "image_thumbnail": image_to_base64(img),
                "image_full": image_to_base64(img),
                "image_full_url": "https://huggingface.co/", # Dummy
                "prompt": "Dummy prompt",
                "category": "Dummy category"
            })
    return evaluation_data

def get_model_names():
    models = ["Qwen2.5-VL", "gemma-3"]
    return models

def get_responses():
    responses = {
        "Qwen2.5-VL": {
            0: "Laws of the Universe - Toro y Moi",
            1: "Smile, have a nice day!",
        },
        "gemma-3": {
            0: "Houdini - Dua Lipa",
            1: "Smile, you're on camera!"
        },
    }
    return responses