"""Download pinned official safetensors, no framework or shared cache."""
import argparse
import concurrent.futures
import hashlib
import json
import pathlib
import os
import shutil
import time
import threading
import urllib.parse
import urllib.request


def fetch_json(url):
    with urllib.request.urlopen(url, timeout=90) as response:
        return json.load(response)


def ranged_download(url, part, total, workers=8):
    """Resume a sequential prefix; journal completed byte ranges before final hashing."""
    journal = part.with_suffix(part.suffix + '.ranges.json')
    if journal.exists():
        state = json.loads(journal.read_text())
        if state['total'] != total:
            raise RuntimeError('Range journal size mismatch')
    else:
        state = {'total': total, 'prefix': part.stat().st_size if part.exists() else 0, 'done': []}
        journal.write_text(json.dumps(state))
    lock = threading.Lock()
    fd = os.open(part, os.O_CREAT | os.O_RDWR, 0o600)
    chunk_size = 64 * 1024 * 1024
    ranges = [(start, min(start + chunk_size, total) - 1)
              for start in range(state['prefix'], total, chunk_size)]

    def one(bounds):
        start, end = bounds
        if start in state['done']:
            return
        for attempt in range(3):
            try:
                req = urllib.request.Request(url, headers={'Range': f'bytes={start}-{end}'})
                with urllib.request.urlopen(req, timeout=180) as response:
                    expected = f'bytes {start}-{end}/{total}'
                    if response.status != 206 or response.headers.get('Content-Range') != expected:
                        raise RuntimeError('Server did not honor exact requested range')
                    offset = start
                    while chunk := response.read(min(4 * 1024 * 1024, end - offset + 1)):
                        written = os.pwrite(fd, chunk, offset)
                        if written != len(chunk):
                            raise RuntimeError('Short positional write')
                        offset += written
                        if offset > end:
                            break
                    if offset != end + 1:
                        raise RuntimeError('Incomplete range')
                with lock:
                    os.fsync(fd)
                    state['done'].append(start)
                    tmp = journal.with_suffix('.tmp')
                    tmp.write_text(json.dumps(state))
                    tmp.replace(journal)
                    done = state['prefix'] + sum(b - a + 1 for a, b in ranges if a in state['done'])
                    print(f'Weight download {done}/{total} bytes ({100 * done / total:.1f}%)', flush=True)
                return
            except Exception:
                if attempt == 2:
                    raise
                time.sleep(2 ** attempt)
    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
            list(pool.map(one, ranges))
    finally:
        os.close(fd)
    journal.unlink()


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--root', type=pathlib.Path, required=True)
    p.add_argument('--repo', default='Qwen/Qwen3.5-2B')
    args = p.parse_args()
    root = args.root.resolve()
    root.mkdir(parents=True, exist_ok=True)
    weights = root / 'weights'
    weights.mkdir(exist_ok=True)
    marker = weights / '.pilot-owned'
    marker.write_text(args.repo)
    old_manifest = root / 'manifest.json'
    pinned = json.loads(old_manifest.read_text())['revision'] if old_manifest.exists() else None
    endpoint = f'https://huggingface.co/api/models/{args.repo}'
    if pinned:
        endpoint += f'/revision/{pinned}'
    info = fetch_json(endpoint + '?blobs=true')
    revision = info['sha']
    files = [f for f in info['siblings'] if f['rfilename'].endswith('.safetensors')
             or f['rfilename'] in ('config.json', 'model.safetensors.index.json', 'README.md')]
    total = sum(f.get('size', 0) for f in files)
    if shutil.disk_usage(root).free < total + 3_000_000_000:
        raise RuntimeError('Insufficient space for checkpoint plus 3 GB margin')
    manifest = {'repo': args.repo, 'revision': revision, 'expected_bytes': total,
                'started_unix': time.time(), 'files': [], 'license': info.get('cardData', {}).get('license')}
    (root / 'source_metadata.json').write_text(json.dumps(info, indent=2))
    for entry in files:
        name = entry['rfilename']
        if pathlib.PurePosixPath(name).name != name:
            raise RuntimeError(f'Unexpected nested path: {name}')
        dest = weights / name
        part = dest.with_suffix(dest.suffix + '.partial')
        digest = hashlib.sha256()
        url = f'https://huggingface.co/{args.repo}/resolve/{revision}/{urllib.parse.quote(name)}?download=true'
        print(f'Downloading {name}, expected {entry.get("size")} bytes', flush=True)
        if entry.get('size', 0) > 128 * 1024 * 1024:
            ranged_download(url, part, entry['size'])
            with part.open('rb') as stream:
                while chunk := stream.read(4 * 1024 * 1024):
                    digest.update(chunk)
        else:
            with urllib.request.urlopen(url, timeout=180) as response, part.open('wb') as out:
                while chunk := response.read(4 * 1024 * 1024):
                    out.write(chunk)
                    digest.update(chunk)
        sha = digest.hexdigest()
        if entry.get('size') and part.stat().st_size != entry['size']:
            raise RuntimeError(f'Size mismatch: {name}')
        expected_sha = entry.get('lfs', {}).get('sha256')
        if expected_sha and sha != expected_sha:
            raise RuntimeError(f'SHA256 mismatch: {name}')
        part.rename(dest)
        manifest['files'].append({'name': name, 'bytes': dest.stat().st_size,
                                  'sha256': sha, 'upstream_sha256': expected_sha})
        (root / 'manifest.json').write_text(json.dumps(manifest, indent=2))
    manifest['completed_unix'] = time.time()
    (root / 'manifest.json').write_text(json.dumps(manifest, indent=2))
    shutil.copy2(weights / 'config.json', root / 'model_config.json')
    print(json.dumps({'revision': revision, 'bytes': sum(f['bytes'] for f in manifest['files'])}), flush=True)


if __name__ == '__main__':
    main()