"""Prepare the reviewed trained checkpoint in a NEW nonproduction directory; never reset caches."""
from pathlib import Path
import hashlib
import json
import shutil
import sys
import time
import urllib.parse
import urllib.request

sys.path.insert(0, str(Path(__file__).resolve().parent.parent/'model'))
from verify import verify_fixture


def prepare(destination):
    destination=Path(destination)
    manifest=json.loads(Path(__file__).with_name('ARTIFACTS.json').read_text())['weights-only']
    revision=manifest['sourceRevision']
    assert revision=='a4f8f3ea906ed274767e9906dbaede7531d660ff'
    assert manifest['license'].startswith('Apache-2.0')
    assert not destination.exists(), 'Use a new isolated directory; existing data is never replaced.'
    assert shutil.disk_usage(destination.parent).free>8*1024**3, 'Require8GiB available for bounded preparation.'
    destination.mkdir(mode=0o700)
    deadline=time.monotonic()+1200
    for entry in manifest['files']:
        name=entry['file']; assert Path(name).name==name and name not in ('.','..')
        url=entry.get('url') or entry['source']
        parsed=urllib.parse.urlsplit(url)
        assert parsed.scheme=='https' and parsed.netloc=='huggingface.co'
        assert parsed.path.startswith('/sshleifer/distilbart-cnn-12-6/resolve/'+revision+'/')
        temporary=destination/(name+'.partial'); count=0; digest=hashlib.sha256()
        with urllib.request.urlopen(url,timeout=60) as response, temporary.open('xb') as stream:
            while block:=response.read(1024*1024):
                count+=len(block)
                assert count<=entry['bytes'] and time.monotonic()<deadline, 'Artifact size/deadline exceeded.'
                digest.update(block); stream.write(block)
        assert count==entry['bytes'] and digest.hexdigest()==entry['sha256'], 'Artifact identity mismatch; do not load partial files.'
        temporary.rename(destination/name)
    verify_fixture(destination,manifest)
    print(json.dumps({'status':'PASS','directory':str(destination.resolve()),'revision':revision,'files':len(manifest['files'])}))


if __name__=='__main__':
    assert len(sys.argv)==2, 'Usage: python prepare.py NEW_NONPRODUCTION_CACHE_DIRECTORY'
    prepare(sys.argv[1])