import json
from pathlib import Path
from typing import Literal

from inspect_ai import Task, task, task_with
from inspect_ai.dataset import MemoryDataset
from inspect_ai.solver import use_tools

from astabench.evals.utils import not_implemented_solver
from astabench.tools import python_session

# Generated via
# ```
# sample_ids = [sample.id for sample in inspect_evals.ds1000.ds1000().dataset]
# random.seed(0)
# random.shuffle(sample_ids)
# validation_ids = sample_ids[:100]
# ```
# Cached explicitly here to avoid any risk of the randomness changing (if e.g.
# the order of ds1000().dataset changes)
SPLIT_IDS_FILE = Path(__file__).parent / "ds1000_splits.json"


@task
def ds1000(split: Literal["validation", "test", "all"] = "all") -> Task:
    """The DS-1000 eval from inspect_evals, using our astabench-oriented
    defaults for the task settings.

    Split can be one of "validation", "test", or "all".  If "all", the
    validation and test sets are combined into a single dataset (the original
    DS-1000 dataset).
    """
    try:
        from inspect_evals.ds1000.ds1000 import ds1000 as inspect_ds1000
    except ImportError:
        raise

    FORMAT_INSTRUCTIONS = "\n\nWrite the remaining python code to append to the program above (but do not repeat the part of the code that is already given in `<code>...</code>`; just write the new code).  Put your answer inside <code> and </code> tags."

    base_task = inspect_ds1000()
    for sample in base_task.dataset:
        assert isinstance(sample.input, str)
        sample.input += FORMAT_INSTRUCTIONS

    # Filter to only include the selected split
    split_ids = json.loads(SPLIT_IDS_FILE.read_text())
    if split == "validation":
        select_ids = set(split_ids["validation_ids"])
        other_ids = set(split_ids["test_ids"])
    elif split == "test":
        select_ids = set(split_ids["test_ids"])
        other_ids = set(split_ids["validation_ids"])
    elif split == "all":
        select_ids = set(split_ids["validation_ids"]) | set(split_ids["test_ids"])
        other_ids = set()
    else:
        raise ValueError(
            f"Unknown split {split}.  Must be one of 'validation', 'test', or 'all'."
        )

    new_dataset = []
    for sample in base_task.dataset:
        if sample.id in select_ids:
            new_dataset.append(sample)
            select_ids.remove(sample.id)
        elif sample.id not in other_ids:
            raise ValueError(
                f"Sample {sample.id} was not found in the validation or test splits."
            )

    if select_ids:
        raise ValueError(f"{split} IDs {select_ids} were not found in the dataset.")
    if len(new_dataset) != (len(base_task.dataset) - len(other_ids)):
        raise ValueError(
            f"Dataset size mismatch (there were samples corresponding to neither validation nor test): {len(new_dataset)} vs {len(base_task.dataset) - len(other_ids)}"
        )

    return task_with(
        base_task,
        name=f"ds1000_{split}",
        dataset=MemoryDataset(new_dataset),
        setup=[use_tools(python_session())],
        solver=not_implemented_solver(),
        sandbox=(
            "docker",
            (
                resources.files("astabench.util.sandbox") / "sandbox_compose.yaml"
            ).as_posix(),
        ),
    )


@task
def ds1000_validation() -> Task:
    """ds1000 but just the validation samples"""
    return ds1000(split="validation")


@task
def ds1000_test() -> Task:
    """ds1000 but just the test samples"""
    return ds1000(split="test")
