import os
import shutil
from pathlib import Path
from typing import Any, Callable, Literal

from huggingface_hub import snapshot_download  # type: ignore
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from platformdirs import user_cache_dir


def gaia_dataset(
    input_prompt: str | None = None,
    subset: Literal[
        "2023_all", "2023_level1", "2023_level2", "2023_level3"
    ] = "2023_all",
    split: Literal["test", "validation"] = "validation",
    filter: Callable[[Sample], bool] = lambda x: True,
) -> Dataset:
    # use user cache dir for dataset
    GAIA_DATASET_LOCATION = (
        Path(user_cache_dir("inspect_evals")) / "gaia_dataset" / "GAIA"
    )

    # download dataset if required
    if not os.path.exists(GAIA_DATASET_LOCATION):
        GAIA_DATASET_LOCATION.mkdir(parents=True, exist_ok=True)
        try:
            snapshot_download(
                repo_id="gaia-benchmark/GAIA",
                repo_type="dataset",
                local_dir=GAIA_DATASET_LOCATION,
            )
        except Exception as ex:
            shutil.rmtree(GAIA_DATASET_LOCATION, True)
            raise ex

    # map record to sample
    def record_to_sample(record: dict[str, Any]) -> Sample:
        # resolve input prompt
        prompt = input_prompt or DEFAULT_INPUT_PROMPT

        # discover files
        id = record["task_id"]
        files_location = GAIA_DATASET_LOCATION / "2023" / split
        files = [file for file in os.listdir(files_location) if str(id) in file]
        if len(files) > 0:
            file = f"The following file is referenced in the question below and you will likely need to use it in order to find the correct answer: /shared_files/{files[0]}"
        else:
            file = ""

        # map fields
        sample = Sample(
            input=prompt.format(file=file, question=record["Question"]),
            target=record["Final answer"],
            id=id,
            metadata={
                "level": record["Level"],
                "Annotator Metadata": record["Annotator Metadata"],
            },
            setup="mkdir -p /shared_files/",
        )

        # provide sample files
        if len(files) > 0:
            sample.files = {
                "/shared_files/" + files[0]: (files_location / files[0]).as_posix()
            }

        return sample

    # read dataset
    dataset = hf_dataset(
        path=GAIA_DATASET_LOCATION.as_posix(),
        name=subset,
        split=split,
        sample_fields=record_to_sample,
        trust=True,  # Trust GAIA's remote code execution during dataset loading
    )

    # apply filter (if any) and return
    return dataset.filter(filter)


DEFAULT_INPUT_PROMPT = """Please answer the question below. You should:

- Return only your answer, which should be a number, or a short phrase with as few words as possible, or a comma separated list of numbers and/or strings.
- If the answer is a number, return only the number without any units unless specified otherwise.
- If the answer is a string, don't include articles, and don't use abbreviations (e.g. for states).
- If the answer is a comma separated list, apply the above rules to each element in the list.

{file}

Here is the question:

{question}"""
