"""Download one official checkpoint per candidate and immutable source snapshots."""
import hashlib
import json
import pathlib
import tarfile
import urllib.request

root = pathlib.Path(__file__).resolve().parent
for name in ('sources', 'weights', 'reports'):
    (root/name).mkdir(exist_ok=True)
manifest = {'sources': {}, 'weights': {}}
path = root/'reports/assets.json'

def fetch(url, dest):
    if dest.exists():
        raise RuntimeError('Refusing to overwrite '+str(dest))
    req = urllib.request.Request(url, headers={'User-Agent':'Contextra-action-eval-prep'})
    digest = hashlib.sha256()
    with urllib.request.urlopen(req, timeout=120) as response, dest.open('xb') as out:
        while chunk := response.read(1024*1024):
            out.write(chunk)
            digest.update(chunk)
    return {'url': url, 'file': str(dest.relative_to(root)),
            'bytes': dest.stat().st_size, 'sha256': digest.hexdigest()}

for repo, branch, name in [('open-mmlab/mmaction2','main','mmaction2'),
                           ('yjh0410/YOWOv2','master','yowov2')]:
    req = urllib.request.Request('https://api.github.com/repos/'+repo+'/commits/'+branch,
                                 headers={'User-Agent':'Contextra-action-eval-prep'})
    with urllib.request.urlopen(req, timeout=30) as response:
        revision = json.load(response)['sha']
    entry = fetch('https://codeload.github.com/'+repo+'/tar.gz/'+revision,
                  root/'sources'/(name+'.tar.gz'))
    entry['revision'] = revision
    with tarfile.open(root/entry['file']) as archive:
        top = archive.getmembers()[0].name.split('/')[0]
        archive.extractall(root/'sources', filter='data')
    entry['directory'] = 'sources/'+top
    manifest['sources'][name] = entry
    path.write_text(json.dumps(manifest, indent=2))
    print('SOURCE',name,revision,flush=True)

base = 'https://download.openmmlab.com/mmaction/v1.0/'
weights = {
 'stgcnpp': base+'skeleton/stgcnpp/stgcnpp_8xb16-joint-u100-80e_ntu60-xsub-keypoint-2d/stgcnpp_8xb16-joint-u100-80e_ntu60-xsub-keypoint-2d_20221228-86e1e77a.pth',
 'posec3d': base+'skeleton/posec3d/slowonly_r50_8xb16-u48-240e_ntu60-xsub-keypoint/slowonly_r50_8xb16-u48-240e_ntu60-xsub-keypoint_20220815-38db104b.pth',
 'ava_slowonly': base+'detection/slowonly/slowonly_kinetics400-pretrained-r50_8xb16-4x16x1-20e_ava21-rgb/slowonly_kinetics400-pretrained-r50_8xb16-4x16x1-20e_ava21-rgb_20220906-953ef5fe.pth',
 'yowov2': 'https://github.com/yjh0410/YOWOv2/releases/download/yowo_v2_weight/yowo_v2_large_ava.pth',
}
for name,url in weights.items():
    try:
        manifest['weights'][name] = fetch(url,root/'weights'/(name+'.pth'))
        print('WEIGHT',name,manifest['weights'][name]['bytes'],flush=True)
    except Exception as e:
        manifest['weights'][name] = {'url':url,'status':'BLOCKED','error_type':type(e).__name__}
        print('BLOCKED',name,type(e).__name__,flush=True)
    path.write_text(json.dumps(manifest, indent=2))