#!/usr/bin/env python3
"""Run on your CLOUD machine: python download-dataset.py --destination /data/tiktok
Downloads ~289.47 GB to that machine, resumes partial files, verifies SHA-256.
Dataset research-use terms: https://huggingface.co/datasets/kuben-developer/tiktok-videos-4b
"""
import argparse,hashlib,json,pathlib,re,shutil,time,urllib.request,urllib.error
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--destination',required=True,type=pathlib.Path)
parser.add_argument('--base',default='https://data.m90.ai')
args=parser.parse_args();args.destination.mkdir(parents=True,exist_ok=True)
with urllib.request.urlopen(urllib.request.Request(args.base.rstrip('/')+'/manifest.json',headers={'User-Agent':'M90-Dataset-Downloader/1.0'}),timeout=60) as response:manifest=json.load(response)
def checksum(path):
    digest=hashlib.sha256()
    with path.open('rb') as stream:
        for chunk in iter(lambda:stream.read(1024*1024),b''):digest.update(chunk)
    return digest.hexdigest()
for file in manifest['files']:
    name=file['name']
    if not re.fullmatch(r'videos-\d{2}\.parquet',name):raise RuntimeError('Unsafe filename in manifest')
    final=args.destination/name;partial=args.destination/(name+'.part')
    if final.exists():
        if final.stat().st_size==file['size'] and checksum(final)==file['sha256']:
            print(name,'already verified',flush=True);continue
        raise RuntimeError(str(final)+' exists but failed verification; move it before retrying')
    for attempt in range(8):
        offset=partial.stat().st_size if partial.exists() else 0
        if offset>file['size']:raise RuntimeError('Partial file is larger than expected')
        if offset==file['size']:break
        if shutil.disk_usage(args.destination).free<file['size']-offset:raise RuntimeError('Not enough space on this machine for the next file')
        request=urllib.request.Request(args.base.rstrip('/')+'/api/files/'+name,headers={'User-Agent':'M90-Dataset-Downloader/1.0',**({'Range':'bytes='+str(offset)+'-'} if offset else {})})
        try:
            with urllib.request.urlopen(request,timeout=180) as response:
                if offset and (response.status!=206 or response.headers.get('Content-Range')!='bytes '+str(offset)+'-'+str(file['size']-1)+'/'+str(file['size'])):
                    raise RuntimeError('Server did not honor resume range; partial file preserved')
                with partial.open('ab' if offset else 'wb') as stream:
                    for chunk in iter(lambda:response.read(1024*1024),b''):stream.write(chunk)
            if partial.stat().st_size==file['size']:break
        except (OSError,urllib.error.URLError) as exc:
            if attempt==7:raise
            print(name,'retrying:',exc,flush=True);time.sleep(min(30,2**attempt))
    if not partial.exists() or partial.stat().st_size!=file['size'] or checksum(partial)!=file['sha256']:
        raise RuntimeError(name+' failed SHA-256 verification; partial file preserved for inspection')
    partial.replace(final);print(name,'verified',flush=True)
print(str(len(manifest['files']))+' files downloaded and verified.')