"""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

# CN mirrors tried in order before/alongside the daemon's configured mirrors.
# Some namespaces (e.g. swebench/*) are blocked by individual CN mirrors, so we
# fall through: daemon default -> 1ms.run -> baidubce -> sjtug.
_CN_MIRROR_FALLBACKS = [
    '{img}',  # daemon default (uses its own registry-mirrors config)
    'docker.1ms.run/{img}',
    'mirror.baidubce.com/{img}',
    'docker.mirrors.sjtug.sjtu.edu.cn/{img}',
    'hub.rat.dev/{img}',
]


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 images_for_samples(samples) -> List[str]:
    """Distinct sandbox images in sample order (dedup, keeps order)."""
    seen: List[str] = []
    for s in samples:
        if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
            seen.append(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:
    """Pull via CN-mirror fallback chain; retag to the canonical name on hit."""
    last_err = None
    for template in _CN_MIRROR_FALLBACKS:
        ref = template.format(img=image)
        r = subprocess.run(['docker', 'pull', ref], capture_output=True, text=True,
                           timeout=3600)
        if r.returncode == 0:
            if ref != image:  # retag the mirrored pull to the canonical name
                subprocess.run(['docker', 'tag', ref, image], check=False)
                subprocess.run(['docker', 'rmi', ref], check=False)
            return
        last_err = (ref, (r.stderr or '').strip()[:150])
    raise RuntimeError(f'all mirrors failed for {image}: {last_err}')