#!/usr/bin/env python3
"""
预拉取 SWE-bench 评测所需的 Docker 镜像，支持断点续传。

用法:
    # 默认从上次进度继续
    python bash/pull_swe_bench_images.py --dataset swe_bench_verified --max-workers 4

    # 强制重新开始（删除进度文件）
    rm -f output/.swe_bench_pull_state.json
    python bash/pull_swe_bench_images.py --dataset swe_bench_verified --max-workers 4

    # 测试前 3 个镜像
    python bash/pull_swe_bench_images.py --dry-run 3 --max-workers 1

进度文件:
    output/.swe_bench_pull_state.json
    记录已拉取/跳过/失败的镜像，Ctrl+C 后下次自动跳过。
"""

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

# 国内网络加速（必须在导入 evalscope/huggingface 之前设置）
os.environ.setdefault('USE_MODELSCOPE_HUB', '1')
os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
os.environ.setdefault('PYTHONUNBUFFERED', '1')

from tqdm import tqdm

# 让 evalscope 可导入
PROJECT_ROOT = Path(__file__).parent.parent.resolve()
sys.path.insert(0, str(PROJECT_ROOT / 'evalscope'))

from evalscope.api.dataset import RemoteDataLoader, FieldSpec


DATASET_IDS = {
    'swe_bench_verified': 'princeton-nlp/SWE-bench_Verified',
    # 'swe_bench_lite': 'princeton-nlp/SWE-bench_Lite',
    # 'swe_bench_verified_mini': 'evalscope/swe-bench-verified-mini',
    # 'swe_bench_pro': 'evalscope/swe-bench-pro',
}

DEFAULT_OUTPUT_DIR = PROJECT_ROOT / 'output'
STATE_FILENAME = '.swe_bench_pull_state.json'


def get_state_path(output_dir: Path) -> Path:
    return Path(output_dir) / STATE_FILENAME


def load_state(output_dir: Path) -> dict:
    """Load previous pull state; return empty state if none exists."""
    state_path = get_state_path(output_dir)
    if not state_path.exists():
        return {'done': [], 'skipped': [], 'failed': []}
    try:
        with state_path.open('r', encoding='utf-8') as f:
            state = json.load(f)
        # Ensure required keys exist
        for key in ('done', 'skipped', 'failed'):
            state.setdefault(key, [])
        return state
    except Exception as e:
        print(f'WARNING: failed to load state file {state_path}: {e}')
        return {'done': [], 'skipped': [], 'failed': []}


def save_state(output_dir: Path, state: dict):
    """Atomically save current pull state."""
    state_path = get_state_path(output_dir)
    state_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = state_path.with_suffix('.tmp')
    try:
        with tmp_path.open('w', encoding='utf-8') as f:
            json.dump(state, f, ensure_ascii=False, indent=2)
        tmp_path.replace(state_path)
    except Exception as e:
        print(f'WARNING: failed to save state file {state_path}: {e}')


def load_swe_samples(dataset_id: str):
    """Load SWE-bench samples and wrap them into the format expected by build_images."""
    print(f'正在加载数据集元数据: {dataset_id} ...')
    loader = RemoteDataLoader(
        data_id_or_path=dataset_id,
        split='test',
        sample_fields=FieldSpec(
            input='problem_statement',
            metadata=[
                'instance_id',
                'repo',
                'base_commit',
                'patch',
                'PASS_TO_PASS',
                'FAIL_TO_PASS',
                'test_patch',
                'version',
                'environment_setup_commit',
                'hints_text',
                'created_at',
            ],
        ),
    )
    return loader.load()


def get_image_name(instance, namespace: str, arch: str) -> str:
    """根据 swebench 规则生成远程镜像名称。"""
    if isinstance(instance, dict):
        metadata = instance
    else:
        metadata = getattr(instance, 'metadata', {}) or {}

    instance_id = metadata.get('instance_id', '')
    if not instance_id and not isinstance(instance, dict):
        instance_id = getattr(instance, 'instance_id', '')

    if not instance_id:
        return None

    image_name = f'{namespace}/sweb.eval.{arch}.{instance_id.lower()}:latest'
    image_name = image_name.replace('__', '_1776_')
    return image_name


def get_local_images() -> set:
    """Return a set of all locally available Docker image names."""
    result = subprocess.run(
        ['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        return set()
    return {line.strip() for line in result.stdout.splitlines() if line.strip()}


_print_lock = threading.Lock()


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


def _stream_pull(image_name: str) -> tuple:
    """Run docker pull and stream output in real time.

    Returns (returncode, last_lines_of_output).
    """
    process = subprocess.Popen(
        ['docker', 'pull', image_name],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    lines = []
    if process.stdout:
        for line in process.stdout:
            line = line.rstrip('\n')
            lines.append(line)
            _log(f'   {line}')
            # Keep only the last 50 lines for error reporting
            if len(lines) > 50:
                lines.pop(0)

    process.wait()
    return process.returncode, '\n'.join(lines)


def pull_one_image(args_tuple, state: dict, output_dir: Path, local_images: set):
    """Pull a single image and update shared state.

    Returns (idx, image_name, status, message).
    """
    idx, total, instance, namespace, arch = args_tuple
    image_name = get_image_name(instance, namespace, arch)
    if not image_name:
        return idx, None, 'skip', 'missing instance_id'

    # Resume: skip if already recorded in state
    if image_name in state['done'] or image_name in state['skipped']:
        return idx, image_name, 'skip', 'already recorded in state'

    # Also check local docker images (in case state was deleted)
    if image_name in local_images:
        state['skipped'].append(image_name)
        save_state(output_dir, state)
        return idx, image_name, 'skip', 'already exists locally'

    _log(f'⬇️  [{idx+1}/{total}] START  {image_name}')

    rc, output = _stream_pull(image_name)

    if rc == 0:
        state['done'].append(image_name)
        save_state(output_dir, state)
        return idx, image_name, 'success', ''
    else:
        msg = output.strip()[-500:] if output.strip() else 'docker pull failed'
        state['failed'].append(image_name)
        save_state(output_dir, state)
        return idx, image_name, 'fail', msg


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Pre-pull SWE-bench Docker images with resume support')
    parser.add_argument(
        '--dataset',
        default='swe_bench_verified',
        choices=list(DATASET_IDS.keys()),
        help='SWE-bench dataset variant',
    )
    parser.add_argument(
        '--max-workers',
        type=int,
        default=4,
        help='Max concurrent docker pulls (default: 4)',
    )
    parser.add_argument(
        '--force-arch',
        default='x86_64',
        choices=['', 'arm64', 'x86_64'],
        help='Image architecture (default: x86_64)',
    )
    parser.add_argument(
        '--dockerhub-username',
        default='swebench',
        help='DockerHub namespace for remote images (default: swebench)',
    )
    parser.add_argument(
        '--output-dir',
        default=str(DEFAULT_OUTPUT_DIR),
        help='Directory to store progress state (default: output/)',
    )
    parser.add_argument(
        '--dry-run',
        type=int,
        default=0,
        help='Only test the first N images',
    )
    parser.add_argument(
        '--retry-failed',
        action='store_true',
        help='Retry images that failed in previous runs',
    )
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    state = load_state(output_dir)

    # Optionally retry previously failed images
    if args.retry_failed:
        failed = state['failed']
        state['failed'] = []
        print(f'🔄 重试 {len(failed)} 个之前失败的镜像')

    dataset_id = DATASET_IDS[args.dataset]
    samples = load_swe_samples(dataset_id)

    if args.dry_run > 0:
        samples = samples[:args.dry_run]
        print(f'⚠️  测试模式：仅处理前 {args.dry_run} 个样本以验证镜像名和网络...')
    else:
        print(f'✅ 成功加载 {len(samples)} 个样本。')

    already_done = len(state['done']) + len(state['skipped'])
    print(f'📂 进度文件: {get_state_path(output_dir)}')
    print(f'   已下载/已存在: {len(state["done"])} | 已跳过: {len(state["skipped"])} | 历史失败: {len(state["failed"])}')

    # Cache local image names once to avoid 500+ `docker images -q` calls.
    print('🔄 正在扫描本地镜像...')
    local_images = get_local_images()
    print(f'   本地已有 {len(local_images)} 个镜像')
    print('-' * 60)

    tasks = [
        (idx, len(samples), sample, args.dockerhub_username, args.force_arch)
        for idx, sample in enumerate(samples)
    ]

    success_count = 0
    skip_count = 0
    fail_count = 0

    with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
        futures = {executor.submit(pull_one_image, t, state, output_dir, local_images): t[0] for t in tasks}
        for future in tqdm(as_completed(futures), total=len(futures), desc='Pulling images', unit='img', mininterval=1.0):
            idx = futures[future]
            try:
                _, image_name, status, msg = future.result()
            except Exception as e:
                image_name, status, msg = None, 'fail', str(e)

            if status == 'success':
                success_count += 1
                _log(f'✅ [{idx+1}/{len(samples)}] DONE   {image_name}')
            elif status == 'skip':
                skip_count += 1
                _log(f'⏭️  [{idx+1}/{len(samples)}] SKIP   {image_name or msg}')
            else:
                fail_count += 1
                _log(f'❌ [{idx+1}/{len(samples)}] FAILED {image_name}')
                if msg:
                    _log(f'   {msg}')

    print('=' * 60)
    print('🎉 预拉取任务结束！')
    print(f'   总计: {len(samples)} | 已存在/跳过: {skip_count} | 本次成功: {success_count} | 本次失败: {fail_count}')
    print(f'   进度文件: {get_state_path(output_dir)}')
    print('=' * 60)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print('\n\n⚠️  用户手动中断。进度已保存，下次运行会自动跳过已完成的镜像。')
        sys.exit(0)