from pathlib import Path
from typing import Any

import requests
from inspect_ai.dataset import Dataset, Sample, json_dataset
from platformdirs import user_cache_dir

CACHE_DIR = Path(user_cache_dir("inspect_evals")) / "scicode" / "data"
EVAL_DIR = Path(__file__).parent


def get_dataset(include_dev_set: bool) -> Dataset:
    """
    Create an Inspect Dataset to run the eval.

    Downloads problems and test data (if they are not already downloaded; reads json to Inspect Dataset.
    """
    if include_dev_set:
        maybe_download_problems_incl_dev()
    else:
        maybe_download_problems_excl_dev()

    maybe_download_test_data()

    def record_to_sample(record: dict[str, Any]) -> Sample:
        return Sample(
            input=record["problem_id"], id=record["problem_id"], metadata=record
        )

    problems_path = (
        CACHE_DIR / "problems_incl_dev.jsonl"
        if include_dev_set
        else CACHE_DIR / "problems_excl_dev.jsonl"
    )

    return json_dataset(problems_path.as_posix(), sample_fields=record_to_sample)


def maybe_download_test_data() -> None:
    """
    Download test data from Google Drive, if it is not yet downloaded.

    Downloads file into local eval directory so it can be copied into the Docker container
    and referenced from the scorer scripts.
    """
    import gdown  # type: ignore

    test_data_path = EVAL_DIR / "test_data.h5"
    if not test_data_path.exists():
        try:
            url = "https://drive.google.com/uc?id=17G_k65N_6yFFZ2O-jQH00Lh6iaw3z-AW"
            gdown.download(url, output=test_data_path.as_posix())
        except Exception as e:
            raise Exception(f"Failed to download test data from {url}: {e}")

    assert test_data_path.exists()


def maybe_download_problems_excl_dev() -> None:
    """Download SciCode problem descriptions from GitHub, if they are not yet downloaded, excluding dev set."""
    problems_path = CACHE_DIR / "problems_excl_dev.jsonl"
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if not problems_path.exists():
        url = "https://raw.githubusercontent.com/scicode-bench/SciCode/69a8cfc/eval/data/problems_all.jsonl"
        response = requests.get(url)
        if response.status_code == 200:
            with open(problems_path, "wb") as f:
                f.write(response.content)
        else:
            raise Exception(f"Failed to download problem set from github at {url}.")

    remove_hardcoded_imports(problems_path)
    assert problems_path.exists()


def maybe_download_problems_incl_dev() -> None:
    """Download SciCode problem descriptions from GitHub, if they are not yet downloaded, including dev set."""
    problems_path = CACHE_DIR / "problems_incl_dev.jsonl"
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if not problems_path.exists():
        for url in (
            "https://raw.githubusercontent.com/scicode-bench/SciCode/69a8cfc/eval/data/problems_all.jsonl",
            "https://raw.githubusercontent.com/scicode-bench/SciCode/69a8cfc/eval/data/problems_dev.jsonl",
        ):
            response = requests.get(url)
            if response.status_code == 200:
                with open(problems_path, "ab") as f:
                    f.write(response.content)
            else:
                raise Exception(f"Failed to download problem set from github at {url}.")

    remove_hardcoded_imports(problems_path)
    assert problems_path.exists()


def remove_hardcoded_imports(path: Path) -> None:
    """
    The problem set has two dependencies on the original SciCode module hardcoded into test cases. These would cause errors in the Inspect implementation, so we remove them.

    The two removed functions are reproduced in test_util.py, which is copied to the testing sandbox. Analogous import statements are added dynamically to test cases, with other dependencies, the scorer module.
    """
    with open(path, "r") as f:
        problems = f.read()

    cleaned_problems = problems.replace(
        "from scicode.compare.cmp import cmp_tuple_or_list", ""
    ).replace("from scicode.compare.cmp import are_dicts_close", "")

    with open(path, "w") as f:
        f.write(cleaned_problems)
