"""Materialize the pinned 1,000 train / 250 test SmolDataEnvs tasks for HF Sandboxes.

Adapted from the FineEnvs env Space's prepare.py (same pinned revisions, same
instruction and grader hash checks), with one change: every task.toml gets

    [environment] docker_image = "docker.io/savatar101/env-data-agent-train:base"

because Harbor's hf-sandbox backend needs a prebuilt image and will not build the
task's environment/Dockerfile. The Dockerfile is a one-line FROM of exactly this
image, so the sandbox runs bit-identical content.

Run once at job startup, then point OPENENV_DATASETS at the output:

    python prepare_tasks.py --output /app/prepared
    export OPENENV_DATASETS=/app/prepared/datasets/train,/app/prepared/datasets/test
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import tomllib
from pathlib import Path

ROOT = Path(__file__).resolve().parent

# All tasks share this one base image (their Dockerfile is `FROM <this>` + WORKDIR).
DOCKER_IMAGE = "docker.io/savatar101/env-data-agent-train:base"

DATASETS = {
    "train": (
        "FineEnvs/SmolDataEnvs-harbor-train",
        "4719635555de1666f374d847baebb500d368e493",
    ),
    "test": (
        "FineEnvs/SmolDataEnvs-harbor-test",
        "b130595b579fc6026ad5762925a9a3b40e72f5bb",
    ),
}
GRADER_SHA256 = "16f906e9248fa8b7e1141939e425e88e59d52d65e98e360bbacef3720a43629f"


def task_rows(split):
    return json.loads((ROOT / "data" / f"{split}.json").read_text())


def prepare(destination, splits=("train", "test")):
    import tomli_w
    from huggingface_hub import snapshot_download

    destination = Path(destination).resolve()
    for split in splits:
        rows = task_rows(split)
        repo, revision = DATASETS[split]
        source = Path(
            snapshot_download(
                repo,
                repo_type="dataset",
                revision=revision,
                allow_patterns=[f"tasks/{r['name']}/*" for r in rows],
                max_workers=16,
            )
        )
        for row in rows:
            folder = source / "tasks" / row["name"]
            instruction = (folder / "instruction.md").read_bytes()
            if hashlib.sha256(instruction).hexdigest() != row["instruction_sha256"]:
                raise ValueError(f"Historical task instruction changed: {split}/{row['name']}")
            grader = (folder / "tests/grader.py").read_bytes()
            if hashlib.sha256(grader).hexdigest() != GRADER_SHA256:
                raise ValueError(f"Expected corrected numeric-only Harbor grader: {folder}")
            target = destination / "datasets" / split / "tasks" / row["name"]
            target.mkdir(parents=True, exist_ok=True)
            for rel in (
                "instruction.md",
                "task.toml",
                "environment/Dockerfile",
                "environment/pull_bucket.py",
                "tests/grader.py",
                "tests/test.sh",
            ):
                dst = target / rel
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(folder / rel, dst)
            spec = tomllib.loads((target / "task.toml").read_text())
            spec["environment"]["memory_mb"] = 4096
            spec["environment"]["cpus"] = 1
            # hf-sandbox runs a prebuilt image instead of building environment/Dockerfile.
            spec["environment"]["docker_image"] = DOCKER_IMAGE
            # Cluster user IDs may not exist inside the sandbox user namespace.
            spec["environment"].setdefault("env", {})["TAR_OPTIONS"] = "--no-same-owner"
            (target / "task.toml").write_text(tomli_w.dumps(spec))
        print(f"{split}: {len(rows)} tasks -> {destination / 'datasets' / split}")
    # Overlap guard, as in the original.
    if set(splits) == {"train", "test"}:
        def names(split):
            return {r["name"] for r in task_rows(split)}
        notebooks = [
            {r.get("notebook") for r in task_rows(s)} - {None} for s in ("train", "test")
        ]
        if names("train") & names("test") or notebooks[0] & notebooks[1]:
            raise ValueError("Train/test overlap")
        print("no train/test overlap")


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--output", type=Path, default=ROOT / "prepared")
    p.add_argument("--splits", nargs="+", default=["train", "test"], choices=["train", "test"])
    a = p.parse_args()
    prepare(a.output, tuple(a.splits))