from typing import Any

from inspect_ai.dataset import Sample
from inspect_ai.model import (
    ChatMessageUser,
    ContentAudio,
    ContentImage,
    ContentReasoning,
    ContentText,
    ContentVideo,
)

from .utils import convert_bytes_to_base64_encoded_image

UNCERTAIN_ANSWER_CHOICE = "Insufficient information to answer the question."


# map records to inspect samples (note that this way the "ideal" target is always choice "A" in the dataset)
# Each subset will shuffle the presentation of choices in multiple_choice()
def record_to_sample_base(record: dict[str, Any]) -> Sample:
    return Sample(
        input=record["question"],
        choices=[record["ideal"]] + record["distractors"] + [UNCERTAIN_ANSWER_CHOICE],
        target="A",
    )


def record_to_sample_suppqa(record: dict[str, Any]) -> Sample:
    return Sample(
        input="Paper title: "
        + record["paper-title"]
        + "\nDOI: "
        + record["source"]
        + "\n"
        + record["question"],
        choices=[record["ideal"]] + record["distractors"] + [UNCERTAIN_ANSWER_CHOICE],
        target="A",
    )


def record_to_sample_figqa(record: dict[str, Any]) -> Sample:
    return Sample(
        input=[
            ChatMessageUser(
                content=[
                    ContentImage(
                        image=convert_bytes_to_base64_encoded_image(
                            "image/jpeg", record["figure"]["bytes"]
                        )
                    ),
                    ContentText(text=record["question"]),
                ]
            )
        ],
        choices=[record["ideal"]] + record["distractors"] + [UNCERTAIN_ANSWER_CHOICE],
        target="A",
    )


def record_to_sample_tableqa(record: dict[str, Any]) -> Sample:
    sample_input_content: list[
        ContentText | ContentReasoning | ContentImage | ContentAudio | ContentVideo
    ] = []
    # can have multiple table images
    for table in record["tables"]:
        sample_input_content.append(
            ContentImage(
                image=convert_bytes_to_base64_encoded_image("image/png", table["bytes"])
            )
        )
    sample_input_content.append(ContentText(text=str(record["question"])))
    return Sample(
        input=[ChatMessageUser(content=sample_input_content)],
        choices=[record["ideal"]] + record["distractors"] + [UNCERTAIN_ANSWER_CHOICE],
        target="A",
    )


def record_to_sample_protocolqa(record: dict[str, Any]) -> Sample:
    return Sample(
        input="Protocol: " + record["protocol"] + "\n\n" + record["question"],
        choices=[record["ideal"]] + record["distractors"] + [UNCERTAIN_ANSWER_CHOICE],
        target="A",
    )
