"""Download pinned configs and checksum-equivalent official weights."""
import argparse
import json
from pathlib import Path
from urllib.request import urlopen

from scripts.integrity import verify_files

REPO = 'Qwen/Qwen-Image-2.1'
REVISION = 'b3179ad355be050328e483a9dfdd9e60cd62adfa'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--source', choices=['modelscope', 'huggingface'], default='modelscope')
    ap.add_argument('--output', default='models/bf16')
    args = ap.parse_args()
    root = Path(args.output)
    root.mkdir(parents=True, exist_ok=True)
    with urlopen(f'https://huggingface.co/api/models/{REPO}/revision/{REVISION}?blobs=true', timeout=60) as r:
        info = json.load(r)
    if info['sha'] != REVISION:
        raise ValueError('Unexpected upstream revision')
    artifacts = Path('artifacts/upstream')
    artifacts.mkdir(parents=True, exist_ok=True)
    (artifacts / 'huggingface-model-info.json').write_text(json.dumps(info, indent=2), encoding='utf-8')
    if args.source == 'huggingface':
        from huggingface_hub import snapshot_download
        snapshot_download(REPO, revision=REVISION, local_dir=str(root), max_workers=4)
    else:
        # ModelScope 1.33 accepts branch/tag revisions, not arbitrary Git commits.
        # We accept files from master only when their SHA256 matches the pinned HF snapshot.
        from modelscope import snapshot_download
        snapshot_download(REPO, revision='master', local_dir=str(root),
                          allow_patterns=['*.safetensors'], max_workers=4)
        for entry in info['siblings']:
            name = entry['rfilename']
            if name.endswith('.safetensors'):
                continue
            with urlopen(f'https://huggingface.co/{REPO}/resolve/{REVISION}/{name}', timeout=60) as r:
                data = r.read()
            path = root / name
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(data)
    rows = [{'path': e['rfilename'], 'size': e['size'], 'sha256': e['lfs']['sha256']}
            for e in info['siblings'] if e['rfilename'].endswith('.safetensors')]
    verified = verify_files(root, rows)
    Path('artifacts/download-verification.json').write_text(
        json.dumps({'revision': REVISION, 'files': verified}, indent=2), encoding='utf-8')
    print('All official weights match the pinned upstream SHA256 manifest.')


if __name__ == '__main__':
    main()