# Adapted from https://github.com/UKGovernmentBEIS/inspect_evals/blob/main/src/inspect_evals/gdm_capabilities/intercode_ctf/dataset.py

import json
import os
import shutil
import zipfile
from pathlib import Path
from textwrap import dedent
from typing import Any

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

USACO_BASE_DIR = Path(user_cache_dir("inspect_evals")) / "usaco"

USACO_DATA_DIR = USACO_BASE_DIR / "data"

# From original paper
PROMPT_PREAMBLE = dedent("""\
            Please reply with a Python 3 solution to the below problem. Make sure to wrap your code in '```python' and '```' Markdown
            delimiters, and include exactly one block of code with the entire solution
            (in the final code step).
            Reason through the problem and:
            1. Restate the problem in plain English
            2. Conceptualize a solution first in plain English
            3. Write a pseudocode solution
            4. Output the final Python solution with your solution steps in comments.
            No outside libraries are allowed.

    """)


def record_to_sample(record: dict[str, Any]) -> Sample:
    # collect files
    problem_id = record["problem_id"]  # <cp_id>_<name>
    cp_id = int(record["cp_id"])  # contest problem ID
    name = record["name"]  # friendly name
    problem_level = record["problem_level"]  # bronze / silver / gold / platinum

    num_tests = int(record["num_tests"])

    runtime_limit = int(record["runtime_limit"])  # max number of seconds
    memory_limit = int(record["memory_limit"])  # max MB (RAM)
    description = record["description"]  # complete prompt text

    # Generate path to test cases
    # - On Windows some filenames contain invalid characters (unzip replaces them with underscores)
    problem_id = (
        problem_id.replace("?", "_").replace(":", "_")
        if os.name == "nt"
        else problem_id
    )
    test_assets_path = USACO_DATA_DIR / "usaco_v3" / "tests" / str(problem_id)

    # Generate prompt
    prompt = PROMPT_PREAMBLE + dedent(f"""\
        [BEGIN PROBLEM]
        {description}
        [END PROBLEM]
    """)

    # Create sample with metadata
    sample = Sample(
        id=cp_id,  # unique identifier for sample
        input=prompt,  # input to be submitted to model
        metadata={
            "name": name,
            "problem_id": problem_id,
            "problem_level": problem_level,  # "bronze" | "silver" | "gold" | "platinum"
            "runtime_limit_secs": runtime_limit,  # Max seconds
            "memory_limit_mbs": memory_limit,  # Max MB (RAM)
            "num_tests": num_tests,
            "test_assets_path": test_assets_path,
        },
    )

    return sample


def read_dataset(
    shuffle: bool = False,
    force_download: bool = False,
    problem_levels: list[str] = ["bronze", "silver", "gold", "platinum"],
    json_basename: str = "usaco_subset307",
) -> Dataset:
    """
    Download and read JSON-based dataset.

    problem_levels: only include problems of levels: "bronze" | "silver" | "gold" | "platinum"
    json_basename: {json_basename}_dict.json contains problem info ("usaco_subset307" or "usaco_v2" for 484 samples)
    """
    # ensure the data is downloaded
    _ensure_data(USACO_DATA_DIR, force_download)

    # Convert JSON from a dict of records to an array of records.
    # This allows us to use the standard `json_dataset` function.
    records_dict_file = USACO_DATA_DIR / f"{json_basename}_dict.json"
    records_array_file = USACO_DATA_DIR / f"{json_basename}_array.json"
    _json_dict_to_array(records_dict_file, records_array_file)

    # Auto-convert JSON records to Samples
    # Option to filter by problem level (e.g. most LLMs cannot solve >silver)
    return json_dataset(
        json_file=str(records_array_file),
        sample_fields=record_to_sample,
        shuffle=shuffle,
    ).filter(
        lambda sample: sample is not None
        and sample.metadata is not None
        and (sample.metadata["problem_level"] in problem_levels)
    )


def _json_dict_to_array(
    records_dict_in_file: Path, records_array_out_file: Path
) -> None:
    """Convert JSON from a dict of records to an array of records, then write to disk."""
    with open(records_dict_in_file, "r") as file_in:
        records_dict = json.load(file_in)

    records_array = list(records_dict.values())
    with open(records_array_out_file, "w") as file_out:
        json.dump(records_array, file_out)


def _ensure_data(data_dir: Path, force_download: bool) -> None:
    zip_url = "https://drive.usercontent.google.com/download?id=1z5ODOJMqyer1QxzYtEUZ2hbAx-7nU8Vi&export=download&authuser=0&confirm=t&uuid=11ea8398-6e91-4a2b-92af-7f6aaf3c2f05&at=AEz70l4HhAdBdGm-QXYS7Y_mknoJ:1741251361013"
    zip_path = str(USACO_BASE_DIR / "usaco_main.zip")
    src_dir = str(USACO_BASE_DIR / "data_copy")

    if data_dir.is_dir() and force_download:
        shutil.rmtree(data_dir)

    data_dir.parent.mkdir(parents=True, exist_ok=True)

    if not data_dir.is_dir():
        if not os.path.exists(zip_path):
            _download_zip(zip_url, zip_path)
        _unzip_file(zip_path, str(USACO_BASE_DIR))
        if os.path.exists(zip_path):
            os.remove(zip_path)

        src = os.path.join(src_dir, "datasets")
        shutil.move(src, data_dir)

        try:
            shutil.rmtree(src_dir)
        except PermissionError:
            pass


def _download_zip(zip_url: str, output_path: str) -> None:
    # Download the zip file
    print(f"Downloading {zip_url}")
    response = requests.get(zip_url)
    if response.status_code == 200:
        with open(output_path, "wb") as file:
            file.write(response.content)
    else:
        raise RuntimeError(
            f"Failed to download source from {zip_url} - {response.status_code}"
        )


def _unzip_file(zip_path: str, extract_to: str) -> None:
    print(f"Unzipping {zip_path}")
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(extract_to)
