#!/usr/bin/env python3
"""Sync the private Hugging Face dataset into the Space local disk."""

from __future__ import annotations

import argparse
import os
import time
from pathlib import Path

from huggingface_hub import snapshot_download


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument('--repo-id', required=True)
    parser.add_argument('--local-dir', type=Path, required=True)
    parser.add_argument('--include', action='append', default=[])
    parser.add_argument('--clean', action='store_true')
    args = parser.parse_args()
    args.local_dir.mkdir(parents=True, exist_ok=True)
    token = os.environ.get('HF_TOKEN')
    if not token:
        raise SystemExit('HF_TOKEN secret is required to sync the private dataset')
    if args.clean:
        for child in args.local_dir.iterdir():
            if child.is_dir():
                import shutil

                shutil.rmtree(child)
            else:
                child.unlink()
    for attempt in range(1, 5):
        try:
            kwargs = {
                'repo_id': args.repo_id,
                'repo_type': 'dataset',
                'local_dir': str(args.local_dir),
                'token': token,
                'max_workers': 2,
            }
            if args.include:
                kwargs['allow_patterns'] = args.include
            snapshot_download(**kwargs)
            return 0
        except Exception as exc:
            if attempt == 4:
                raise
            wait_seconds = 20 * attempt
            print(f'dataset sync failed on attempt {attempt}: {exc}', flush=True)
            print(f'retrying in {wait_seconds}s', flush=True)
            time.sleep(wait_seconds)
    return 1


if __name__ == '__main__':
    raise SystemExit(main())