import functools
from typing import Any

from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.model import ChatMessage, ChatMessageUser


def record_to_sample_with_shot_prompts(
    record: dict[str, Any], shot_prompts: list[ChatMessage] | None = None
) -> Sample:
    # These messages need to be added here since any messages added as part of the solvers
    # go after the task message.
    sample_inputs: list[ChatMessage] = [ChatMessageUser(content=record["task"])]
    if shot_prompts is not None:
        sample_inputs = shot_prompts + sample_inputs

    return Sample(
        id=record["id"],
        input=sample_inputs,
        target=record["answer"],
        metadata={
            "explanation": record["explanation"],
            "difficulty": record["difficulty"],
        },
    )


def assistant_bench_dataset(
    shot_prompts: list[ChatMessage] | None = None,
    shuffle: bool = True,
    seed: int | None = None,
) -> Dataset:
    record_to_sample = functools.partial(
        record_to_sample_with_shot_prompts, shot_prompts=shot_prompts
    )
    return hf_dataset(
        path="AssistantBench/AssistantBench",
        sample_fields=record_to_sample,
        split="validation",
        shuffle=shuffle,
        seed=seed,
    )
