"""Run on 2033lab04ztr: mirror download with per-file size/SHA256 checks."""
import concurrent.futures
import hashlib
import json
import pathlib
import subprocess

ROOT = pathlib.Path('/mnt/58462F4C462F29E8/xiaobailing-turn-detection-experiments')
REPO = 'Soul-AILab/SoulX-Duplug-0.6B'
REVISION = '61701c4ab8193cc1ee2220d3848872ac6c720142'
TARGET = ROOT / 'SoulX-Duplug/pretrained_models'


def main():
    raw = subprocess.check_output(['curl', '-fLsS', '--retry', '3', '--max-time', '60',
        f'https://hf-mirror.com/api/models/{REPO}/revision/{REVISION}?blobs=true'])
    metadata = json.loads(raw)
    assert metadata['sha'] == REVISION
    (ROOT / 'results/model-metadata.json').write_bytes(raw)
    TARGET.mkdir(parents=True, exist_ok=True)
    ms_raw = subprocess.check_output(['curl', '-fLsS', '--retry', '3', '--max-time', '60',
        f'https://modelscope.cn/api/v1/models/{REPO}/repo/files?Revision=master&Recursive=true'])
    ms_data = json.loads(ms_raw)
    assert ms_data['Code'] == 200
    ms_files = {x['Path']: x for x in ms_data['Data']['Files'] if x['Type'] != 'tree'}
    (ROOT / 'results/modelscope-metadata.json').write_bytes(ms_raw)

    def download(item):
        name, size = item['rfilename'], item['size']
        path = TARGET / name
        expected = item.get('lfs', {}).get('sha256')
        def valid():
            if not path.is_file() or path.stat().st_size != size:
                return False
            if expected:
                with path.open('rb') as f:
                    actual = hashlib.file_digest(f, 'sha256').hexdigest()
                return actual == expected
            return True
        if not valid():
            path.parent.mkdir(parents=True, exist_ok=True)
            tmp = path.with_name(path.name + '.download')
            url = f'https://hf-mirror.com/{REPO}/resolve/{REVISION}/{name}?download=true'
            if name in ms_files and size > 1000000:
                ms = ms_files[name]
                assert ms['Size'] == size
                if expected:
                    assert ms['Sha256'] == expected, f'Mirror mismatch: {name}'
                else:
                    expected = ms['Sha256']
                url = f'https://modelscope.cn/api/v1/models/{REPO}/repo?Revision=master&FilePath={name}'
            subprocess.run(['curl', '-fL', '--retry', '5', '--retry-delay', '3',
                '--connect-timeout', '15', '--speed-limit', '1024', '--speed-time', '90',
                '-sS', '-o', str(tmp), url], check=True)
            tmp.replace(path)
            if not valid():
                raise ValueError(f'Invalid downloaded file: {name}')
        print(f'OK {name} {size} bytes', flush=True)
        return dict(file=name, bytes=size, lfs_sha256=expected)

    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
        checked = list(pool.map(download, metadata['siblings']))
    (ROOT / 'results/model-download.json').write_text(json.dumps(dict(
        repo=REPO, revision=REVISION, files=checked), indent=2) + '\n')
    print('DOWNLOAD_COMPLETE', flush=True)


if __name__ == '__main__':
    main()