#!/usr/bin/env python3
"""Download a pinned official ModelScope snapshot; never replace mismatched files."""
import hashlib
import json
import os
from pathlib import Path
import sys
import urllib.parse
import urllib.request

MODEL = 'Qwen/Qwen3-ASR-1.7B'
REVISION = 'a04930dbe5419bfee073f7cade734f572689a3a8'
ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else '/home/www/qwen-vllm/asr/models/Qwen3-ASR-1.7B')
API = f'https://modelscope.cn/api/v1/models/{MODEL}/repo'

def digest(path):
    h = hashlib.sha256()
    with path.open('rb') as stream:
        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b''):
            h.update(chunk)
    return h.hexdigest()

ROOT.mkdir(parents=True, exist_ok=True)
url = API + '/files?' + urllib.parse.urlencode({'Revision': REVISION, 'Recursive': 'true'})
with urllib.request.urlopen(url, timeout=60) as response:
    listing = json.load(response)
assert listing['Code'] == 200 and listing['Success'], listing
files = [x for x in listing['Data']['Files'] if x['Type'] == 'blob' and x['Path'] != '.gitattributes']
manifest = {'provider': 'ModelScope', 'model': MODEL, 'revision': REVISION, 'api': url, 'files': files}
manifest_path = ROOT / 'DOWNLOAD_MANIFEST.json'
if manifest_path.exists():
    old = json.loads(manifest_path.read_text())
    assert old == manifest, 'Existing manifest differs; preserve and abort'
else:
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + '\n')
for row in files:
    relative = Path(row['Path'])
    assert not relative.is_absolute() and '..' not in relative.parts
    path = ROOT / relative
    if path.exists():
        assert path.stat().st_size == row['Size'] and digest(path) == row['Sha256'], f'Existing file mismatch: {path}'
        print(json.dumps({'event': 'download_reused', 'path': str(path), 'sha256': row['Sha256'], 'bytes': row['Size']}), flush=True)
        continue
    path.parent.mkdir(parents=True, exist_ok=True)
    partial = path.with_name(path.name + '.partial')
    # A partial file is only a disposable incomplete transfer created by this script.
    target = API + '?' + urllib.parse.urlencode({'Revision': REVISION, 'FilePath': row['Path']})
    print(json.dumps({'event': 'download_started', 'file': row['Path'], 'bytes': row['Size'], 'revision': REVISION}), flush=True)
    h = hashlib.sha256()
    size = 0
    with urllib.request.urlopen(target, timeout=120) as source, partial.open('wb') as output:
        while True:
            chunk = source.read(8 * 1024 * 1024)
            if not chunk:
                break
            output.write(chunk)
            h.update(chunk)
            size += len(chunk)
    assert size == row['Size'] and h.hexdigest() == row['Sha256'], f'Download integrity mismatch: {path}'
    os.replace(partial, path)
    print(json.dumps({'event': 'download_verified', 'path': str(path), 'sha256': h.hexdigest(), 'bytes': size}), flush=True)
config = json.loads((ROOT / 'config.json').read_text())
assert config['architectures'] == ['Qwen3ASRForConditionalGeneration'], config['architectures']
weights = json.loads((ROOT / 'model.safetensors.index.json').read_text())
assert set(weights['weight_map'].values()) <= {x['Path'] for x in files}
# Public model weights must remain readable through a read-only bind mount with cap_drop=ALL.
# This does not change the private service .env or acceptance directory permissions.
ROOT.chmod(0o755)
manifest_path.chmod(0o644)
for row in files:
    (ROOT / row['Path']).chmod(0o644)
print(json.dumps({'event': 'model_ready', 'model': MODEL, 'revision': REVISION, 'manifest_sha256': digest(manifest_path), 'total_bytes': sum(x['Size'] for x in files), 'architectures': config['architectures'], 'model_type': config.get('model_type'), 'dtype': config.get('torch_dtype', config.get('dtype')), 'weight_metadata': weights['metadata']}), flush=True)