#!/usr/bin/env python3
from __future__ import annotations

import argparse
import docker
from pathlib import Path

from swebench.harness.utils import load_swebench_dataset
from swebench.harness.docker_build import build_env_images


def main():
    p = argparse.ArgumentParser(description="Build only the SWE-bench environment images for given instance IDs.")
    p.add_argument("--dataset_name", type=str, default="SWE-bench/SWE-bench")
    p.add_argument("--split", type=str, default="test")
    p.add_argument("--instance_ids", nargs="+", required=True, help="Space-separated instance IDs")
    p.add_argument("--max_workers", type=int, default=1, help="Workers for env image build (use 1 to reduce RAM)")
    p.add_argument("--force_rebuild", action="store_true", help="Force rebuild env images even if cached")
    args = p.parse_args()

    # Load only the requested instances
    dataset = load_swebench_dataset(args.dataset_name, args.split, instance_ids=args.instance_ids)
    if not dataset:
        print("No matching instances in dataset.")
        return

    client = docker.from_env()
    print(f"Building env images for {len(dataset)} instance(s) with max_workers={args.max_workers}...")
    build_env_images(client, dataset, force_rebuild=args.force_rebuild, max_workers=args.max_workers)
    print("Done.")


if __name__ == "__main__":
    main()
