"""Image prefetch: parallel docker pull for execution environments.

SWE-bench Verified declares ~500 per-instance images (sweb.eval.x86_64.*).
Pulling them lazily during an eval run would stall it serially; prefetch
pulls them ahead of time with bounded concurrency and progress reporting.

    from evalharness.sandbox.prefetch import prefetch_images
    done = prefetch_images(['sweb.eval.x86_64.django__django-12345', ...], workers=8)
    # CLI: evalharness sandbox prefetch swe_bench_verified --workers 8 --limit 50
"""

import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Iterable, List

from ..data.dataset import Dataset


def images_for_dataset(ds: Dataset, limit: int = 0) -> List[str]:
    """Collect distinct sandbox images declared by a dataset's samples."""
    seen: List[str] = []
    add = seen.append
    for i, s in enumerate(ds):
        if limit and i >= limit:
            break
        if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
            add(s.sandbox.image)
    return seen


def local_images() -> set:
    out = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
                         capture_output=True, text=True)
    return {line for line in out.stdout.splitlines() if line}


def prefetch_images(images: Iterable[str], workers: int = 8) -> List[str]:
    """Pull images with bounded concurrency; skip ones already local.

    Returns the list newly pulled. Failures are reported and skipped (one
    missing instance must not block the rest).
    """
    have = local_images()
    todo = [img for img in images if img not in have]
    if not todo:
        print(f'prefetch: all {len(images)} images already local')
        return []
    print(f'prefetch: {len(todo)} to pull ({len(images) - len(todo)} already local), '
          f'workers={workers}')
    pulled: List[str] = []
    failed: List[str] = []
    t0 = time.time()
    done = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futs = {pool.submit(_pull_one, img): img for img in todo}
        for fut in as_completed(futs):
            img = futs[fut]
            done += 1
            try:
                fut.result()
                pulled.append(img)
            except Exception as e:
                failed.append(img)
                print(f'  FAIL {img}: {str(e)[:120]}', flush=True)
            if done % 10 == 0 or done == len(todo):
                rate = done / max(time.time() - t0, 1e-6)
                print(f'  [{done}/{len(todo)}] {rate:.2f} imgs/s, '
                      f'{len(failed)} failed', flush=True)
    print(f'prefetch done: {len(pulled)} pulled, {len(failed)} failed '
          f'in {time.time() - t0:.0f}s')
    return pulled


def _pull_one(image: str) -> None:
    r = subprocess.run(['docker', 'pull', image], capture_output=True, text=True, timeout=3600)
    if r.returncode != 0:
        raise RuntimeError(r.stderr.strip()[:200])