import os
from functools import cache
from typing import Any

from datasets import Dataset  # type: ignore
from huggingface_hub import snapshot_download  # type: ignore
from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.model import ChatMessage, ChatMessageUser, ContentImage, ContentText
from inspect_ai.scorer import choice
from inspect_ai.solver import multiple_choice

HUGGINGFACE_REPO_ID = "craigwu/vstar_bench"


@task
def vstar_bench_attribute_recognition() -> Task:
    """Evaluates model on V*Bench direct_attributes questions."""
    return create_task(category="direct_attributes")


@task
def vstar_bench_spatial_relationship_reasoning() -> Task:
    """Evaluates model on V*Bench relative_position questions."""
    return create_task(category="relative_position")


def create_task(category: str | None) -> Task:
    return Task(
        dataset=(load_dataset().filter(lambda x: x.metadata["category"] == category)),
        solver=multiple_choice(),
        scorer=choice(),
    )


def load_dataset() -> Dataset:
    """Load the V-Star-Bench dataset from HuggingFace and converts it to Inspect dataset."""
    return hf_dataset(
        path=HUGGINGFACE_REPO_ID, split="test", sample_fields=record_to_sample
    )


def record_to_sample(record: dict[str, Any]) -> Sample:
    """Convert a hf record to a sample."""
    text, options = extract_choices(record["text"])
    input: list[ChatMessage] = [
        ChatMessageUser(
            content=[
                ContentText(text=text),
                ContentImage(
                    type="image",
                    image=os.path.join(images_path(), record["image"]),
                ),
            ],
        ),
    ]
    return Sample(
        id=record["question_id"],
        input=input,
        choices=options,
        target=record["label"],
        metadata={"category": record["category"]},
    )


def extract_choices(text: str) -> tuple[str, list[str]]:
    """Extracts the question text and answer choices from the text.

    Text is expected to be in the following format:
    Is the red backpack on the left or right side of the clock?
    (A) right
    (B) left
    Answer with the option's letter from the given choices directly.
    """
    text = text.replace(
        "\nAnswer with the option's letter from the given choices directly.", ""
    )
    lines = text.split("\n")
    question = lines[0]
    choices = [line[4:] for line in lines[1:]]
    assert len(choices) in [2, 4]
    return question, choices


@cache
def images_path() -> str:
    """Downloads dataset snapshot and returns its local path."""
    return snapshot_download(repo_id=HUGGINGFACE_REPO_ID, repo_type="dataset")
