"""ModelScope-hosted sandbox images: per-image tar upload + download.

registry.modelscope.cn does not accept personal docker pushes, so images
ship as ONE docker-save tar per image inside a ModelScope MODEL repo
(file hosting with a stable CN CDN -- measured ~12MB/s anonymous; the
per-task granularity means running 3 instances downloads exactly 3 tars).
This is the only working network source for the swebench/* namespace --
public CN mirrors (daocloud/1ms/...) 403 it.

NB: it must be a MODEL repo, not a dataset repo: upload_file defaults to
repo_type='model' and silently CREATES one when given a dataset repo_id
(the first upload landed in a phantom model repo while the dataset repo
stayed empty -- confusing 404s on pull).

Upload (resumable; skips tars already on the remote):
    EVALHARNESS_MS_TOKEN=ms-... python -m evalharness.sandbox.ms_images \\
        upload --repo SoraAmami/swebench-images [--limit N]

Download side is wired into ImageService._pull_one (first source, before
the docker-hub mirror chain) whenever EVALHARNESS_MS_IMAGE_REPO is set
(defaults on for this deployment; empty string disables).
"""

import os
import subprocess
import sys
import tempfile
import time
from typing import List, Optional, Set

# a private-repo token for downloads (public repos work without it)
_MS_TOKEN = os.environ.get('EVALHARNESS_MS_TOKEN', '').strip()
_MS_ENDPOINT = os.environ.get('EVALHARNESS_MS_ENDPOINT', 'https://modelscope.cn')
# big enough for a ~4GB uncompressed save; / has less headroom than /data1
_TMP_ROOT = os.environ.get('EVALHARNESS_MS_TMP', '/data1/sora/temp/ms_images')
_STATE = os.path.join(_TMP_ROOT, 'uploaded.txt')   # resume manifest

# image prefixes this deployment ships (swe per-instance + code benches)
_DEFAULT_PREFIXES = ('swebench/',)


def _docker(args, timeout=1800):
    return subprocess.run(['docker'] + args, capture_output=True, text=True,
                          timeout=timeout)


def file_for(image: str) -> str:
    """Image ref -> dataset filename: last path segment minus tag.

    swebench/sweb.eval.x86_64.django_1776_django-13964:latest
        -> sweb.eval.x86_64.django_1776_django-13964.tar.gz
    """
    base = image.rsplit('/', 1)[-1]
    base = base.split(':')[0]
    return f'{base}.tar.gz'


def image_from_file(fname: str) -> Optional[str]:
    """Inverse of file_for for the swe namespace (namespace is not
    recoverable in general; callers here only ship swebench/*)."""
    if not fname.endswith('.tar.gz'):
        return None
    return f"swebench/{fname[:-len('.tar.gz')]}:latest"


def _local_images() -> Set[str]:
    r = _docker(['images', '--format', '{{.Repository}}:{{.Tag}}'], timeout=60)
    return set(r.stdout.split()) if r.returncode == 0 else set()


def _remote_files(repo: str) -> Set[str]:
    """Files currently in the MODEL repo (resume manifest source of truth)."""
    from modelscope.hub.api import HubApi

    api = HubApi()
    if _MS_TOKEN:
        api.login(_MS_TOKEN)
    files: Set[str] = set()
    for e in api.get_model_files(repo, recursive=True):
        p = (e.get('Path') or e.get('Name') or '').lstrip('/')
        if p.endswith('.tar.gz'):
            files.add(p)
    return files


def _upload_one(api, repo: str, image: str, tmp_dir: str) -> bool:
    tar = os.path.join(tmp_dir, file_for(image))
    os.makedirs(tmp_dir, exist_ok=True)
    # gzip -1: ~2x faster than default for ~10% more bytes -- upload wall
    # time is dominated by docker save IO, not the extra size
    r = subprocess.run(f'docker save {image} | gzip -1 > {tar}',
                       shell=True, timeout=3600)
    if r.returncode != 0 or not os.path.exists(tar):
        print(f'  save failed: {image}', flush=True)
        return False
    try:
        api.upload_file(repo_id=repo, path_or_fileobj=tar,
                        path_in_repo=file_for(image),
                        repo_type='model', token=_MS_TOKEN or None)
        with open(_STATE, 'a') as f:
            f.write(file_for(image) + '\n')
        gb = os.path.getsize(tar) / 1e9
        print(f'  uploaded {file_for(image)} ({gb:.2f} GB)', flush=True)
        return True
    except Exception as e:
        print(f'  upload failed {image}: {type(e).__name__}: {str(e)[:120]}',
              flush=True)
        return False
    finally:
        try:
            os.unlink(tar)
        except OSError:
            pass


def upload(repo: str, limit: int = 0, prefixes=None) -> None:
    """Ship every local matching image to the MODEL repo. Resumable:
    already-remote tars are skipped (state file + remote listing)."""
    from modelscope.hub.api import HubApi

    os.makedirs(_TMP_ROOT, exist_ok=True)
    api = HubApi()
    if _MS_TOKEN:
        api.login(_MS_TOKEN)
    try:
        api.get_model(repo)
    except Exception:
        print(f'creating model repo {repo} ...', flush=True)
        api.create_model(model_id=repo)
    done = set()
    try:
        done = {l.strip() for l in open(_STATE) if l.strip()}
    except OSError:
        pass
    print(f'remote listing {repo} ...', flush=True)
    try:
        done |= _remote_files(repo)
    except Exception as e:
        print(f'  remote listing failed ({e}); relying on local state',
              flush=True)
    prefixes = prefixes or _DEFAULT_PREFIXES
    imgs = sorted(i for i in _local_images() if i.startswith(prefixes))
    if limit:
        imgs = imgs[:limit]
    todo = [i for i in imgs if file_for(i) not in done]
    print(f'{len(imgs)} local images, {len(todo)} to upload '
          f'({len(imgs) - len(todo)} already remote)', flush=True)
    ok = 0
    for n, img in enumerate(todo, 1):
        print(f'[{n}/{len(todo)}] {img}', flush=True)
        for attempt in range(3):
            if _upload_one(api, repo, img, _TMP_ROOT):
                ok += 1
                break
            time.sleep(10 * (attempt + 1))     # hub hiccups: backoff + retry
    print(f'done: {ok}/{len(todo)} uploaded', flush=True)


def ms_pull(image: str, repo: str) -> bool:
    """Fetch one image tar from the dataset repo and docker-load it.
    Returns True iff the image is local afterwards."""
    import urllib.request

    fname = file_for(image)
    url = (f'{_MS_ENDPOINT}/api/v1/models/{repo}/repo'
           f'?Revision=master&FilePath={fname}')
    hdrs = {'User-Agent': 'evalharness'}
    if _MS_TOKEN:
        hdrs['Authorization'] = f'Bearer {_MS_TOKEN}'
    os.makedirs(_TMP_ROOT, exist_ok=True)
    fd, tar = tempfile.mkstemp(suffix='.tar.gz', dir=_TMP_ROOT)
    os.close(fd)
    try:
        req = urllib.request.Request(url, headers=hdrs)
        t0 = time.time()
        with urllib.request.urlopen(req, timeout=600) as r, open(tar, 'wb') as f:
            while True:
                chunk = r.read(1 << 20)
                if not chunk:
                    break
                f.write(chunk)
        r = _docker(['load', '-qi', tar], timeout=1800)
        if r.returncode == 0 and image in _local_images():
            gb = os.path.getsize(tar) / 1e9
            print(f'· ms-images: loaded {image.split("/")[-1][:50]} '
                  f'({gb:.2f} GB in {time.time() - t0:.0f}s)', flush=True)
            return True
        print(f'· ms-images: load failed for {image}: '
              f'{(r.stderr or "")[:120]}', flush=True)
        return False
    except Exception as e:
        # 404 = this image was never uploaded; anything else = transient
        code = getattr(e, 'code', None)
        if code != 404:
            print(f'· ms-images: fetch {fname}: {type(e).__name__}: '
                  f'{str(e)[:100]}', flush=True)
        return False
    finally:
        try:
            os.unlink(tar)
        except OSError:
            pass


if __name__ == '__main__':
    import argparse

    p = argparse.ArgumentParser()
    p.add_argument('cmd', choices=['upload'])
    p.add_argument('--repo', default='SoraAmami/swebench-images')
    p.add_argument('--limit', type=int, default=0)
    a = p.parse_args()
    if a.cmd == 'upload':
        upload(a.repo, limit=a.limit)