#!/usr/bin/env python3
"""SWE-bench 逐题镜像预拉取（断点续传、并发、失败重试）。

移植自 es 的 pull_swe_bench_images.py，去掉 evalscope 依赖（用我们自己的
数据集注册表），加进 sandbox 子命令：

    evalharness sandbox pull swe_bench_verified --max-workers 4
    evalharness sandbox pull swe_bench_verified --dry-run 3      # 只试前 3 个
    evalharness sandbox pull swe_bench_verified --retry-failed

进度文件: <cache-dir>/swe_pull_state.json（done/failed/skipped），
Ctrl+C 后重跑自动跳过已完成的。
"""

import json
import subprocess
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

_print_lock = threading.Lock()


def _log(msg):
    with _print_lock:
        print(msg, flush=True)


def swe_image_names(dataset_name: str):
    """数据集 → 官方逐题镜像名列表（与 dataset 插件的命名完全一致）。"""
    from evalharness import get_dataset

    ds = get_dataset(dataset_name)
    ds.materialize()
    names = []
    for s in ds:
        if s.sandbox and s.sandbox.image:
            names.append(s.sandbox.image)
    return names


def _state_path(cache_dir: Path) -> Path:
    return cache_dir / 'swe_pull_state.json'


def _load_state(cache_dir: Path) -> dict:
    p = _state_path(cache_dir)
    if not p.exists():
        return {'done': [], 'failed': [], 'skipped': []}
    try:
        state = json.load(open(p, encoding='utf-8'))
        for k in ('done', 'failed', 'skipped'):
            state.setdefault(k, [])
        return state
    except Exception:
        return {'done': [], 'failed': [], 'skipped': []}


def _save_state(cache_dir: Path, state: dict):
    p = _state_path(cache_dir)
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_suffix('.tmp')
    tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding='utf-8')
    tmp.replace(p)


def _pull_one(image: str) -> tuple:
    """docker pull 单个镜像，实时透传输出。返回 (returncode, 尾部输出)。"""
    proc = subprocess.Popen(['docker', 'pull', image],
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            text=True, bufsize=1)
    lines = []
    for line in proc.stdout:
        line = line.rstrip('\n')
        lines.append(line)
        _log(f'   {line}')
        if len(lines) > 50:
            lines.pop(0)
    proc.wait()
    return proc.returncode, '\n'.join(lines)


def pull_many(names, cache_dir, max_workers: int = 4,
              verbose: bool = True) -> int:
    """批量拉取一组镜像名（断点续传复用 <cache-dir>/swe_pull_state.json）。

    返回失败数。给判分预检用：多镜像 bench（swe 500 个）逐个 fail-fast
    会把整个 bench 卡死在第一个失败上；这里并发拉、不中断、返回失败数。"""
    cdir = Path(cache_dir)
    state = _load_state(cdir)
    local = set()
    r = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
                       capture_output=True, text=True)
    if r.returncode == 0:
        local = {l.strip() for l in r.stdout.splitlines() if l.strip()}
    todo = [nm for nm in names
            if nm not in state['done'] and nm not in local]
    if verbose:
        _log(f'· {len(names)} 个沙箱镜像: 本地/已完成 {len(names) - len(todo)}, '
             f'待拉 {len(todo)} (并发 {max_workers}, 进度存 {_state_path(cdir)})')

    def work(nm):
        rc, out = _pull_one(nm)
        if rc == 0:
            state['done'].append(nm)
        else:
            state['failed'].append(nm)
        _save_state(cdir, state)
        return nm, rc

    fails = 0
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        for i, (nm, rc) in enumerate(
                pool.map(work, todo), 1):
            if rc:
                fails += 1
            if verbose and (i % 10 == 0 or i == len(todo)):
                _log(f'  镜像进度 [{i}/{len(todo)}] 失败 {fails}')
    return fails


def run_pull(dataset: str, cache_dir: str, max_workers: int = 4,
             dry_run: int = 0, retry_failed: bool = False):
    cdir = Path(cache_dir)
    state = _load_state(cdir)
    if retry_failed:
        n = len(state['failed'])
        state['failed'] = []
        _log(f'🔄 重试 {n} 个之前失败的镜像')

    names = swe_image_names(dataset)
    if dry_run:
        names = names[:dry_run]
    local = set()
    r = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
                       capture_output=True, text=True)
    if r.returncode == 0:
        local = {l.strip() for l in r.stdout.splitlines() if l.strip()}

    todo = []
    for nm in names:
        if nm in state['done'] or nm in state['skipped']:
            continue
        if nm in local:
            state['skipped'].append(nm)
            _save_state(cdir, state)
            continue
        todo.append(nm)
    _log(f'共 {len(names)} 个镜像: 已完成/本地 {len(names) - len(todo)}, '
         f'待拉 {len(todo)}, 并发 {max_workers}')

    def work(nm):
        _log(f'⬇ START {nm}')
        rc, out = _pull_one(nm)
        if rc == 0:
            state['done'].append(nm)
        else:
            state['failed'].append(nm)
        _save_state(cdir, state)
        return nm, rc, out[-300:] if rc else ''

    fails = []
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futs = {pool.submit(work, nm): nm for nm in todo}
        for i, fut in enumerate(as_completed(futs), 1):
            nm, rc, out = fut.result()
            mark = '✓' if rc == 0 else '✗'
            _log(f'[{i}/{len(todo)}] {mark} {nm}')
            if rc:
                fails.append((nm, out))

    _log(f'\n完成: {len(state["done"])} 成功, {len(state["skipped"])} 跳过, '
         f'{len(state["failed"])} 失败 (进度存 {_state_path(cdir)})')
    for nm, out in fails[:5]:
        _log(f'  失败样例 {nm}: {out[:150]}')
    return 0 if not fails else 1


def main(argv) -> int:
    import argparse

    p = argparse.ArgumentParser(description='Pre-pull SWE-bench per-instance '
                                            'images with resume')
    p.add_argument('dataset', nargs='?', default='swe_bench_verified')
    p.add_argument('--max-workers', type=int, default=4)
    p.add_argument('--dry-run', type=int, default=0,
                   help='only pull the first N images')
    p.add_argument('--retry-failed', action='store_true')
    p.add_argument('--cache-dir', default='')
    a = p.parse_args(argv)

    cache = a.cache_dir
    if not cache:
        import os
        cache = os.environ.get('EVALHARNESS_CACHE',
                               str(Path.home() / '.cache/evalharness'))
    return run_pull(a.dataset, cache, a.max_workers, a.dry_run, a.retry_failed)


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))