"""Prepare project-local Hugging Face cache metadata and resumable downloads."""
import json
import subprocess
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
CACHE = ROOT / '.hf_cache' / 'hub'

def fetch(url, target):
    target.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(['curl', '-fsSL', '--retry', '3', '--retry-all-errors', '--max-time', '45', url, '-o', str(target)], check=True)

manifest = []
for repo, files, aliases in [
    ('google-bert/bert-base-uncased', ['config.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.txt', 'model.safetensors'], ['bert-base-uncased']),
    ('lpiccinelli/unidepth-v2-vitl14', ['config.json', 'model.safetensors'], []),
    ('Rain729/Prior-Depth-Anything', ['depth_anything_v2_vitb.pth', 'prior_depth_anything_vitb.pth'], []),
]:
    meta = ROOT / 'output' / (repo.replace('/', '_') + '_metadata.json')
    fetch('https://hf-mirror.com/api/models/' + repo, meta)
    revision = json.loads(meta.read_text())['sha']
    cache = CACHE / ('models--' + repo.replace('/', '--'))
    snapshot = cache / 'snapshots' / revision
    snapshot.mkdir(parents=True, exist_ok=True)
    (cache / 'refs').mkdir(exist_ok=True)
    (cache / 'refs' / 'main').write_text(revision)
    for name in files:
        target = snapshot / name
        url = 'https://hf-mirror.com/' + repo + '/resolve/' + revision + '/' + name
        if repo.startswith('lpiccinelli/'):
            local = ROOT / 'weights' / 'unidepth-v2-vitl14' / name
            if not target.exists() and not target.is_symlink():
                target.symlink_to(local)
        elif name.endswith(('.json', '.txt')):
            fetch(url, target)
        else:
            manifest.extend([url, '  dir=' + str(snapshot), '  out=' + name])
    for alias in aliases:
        alias_path = CACHE / ('models--' + alias.replace('/', '--'))
        if not alias_path.exists():
            alias_path.symlink_to(cache, target_is_directory=True)
(ROOT / 'output' / 'vipe_weights_downloads.txt').write_text('\n'.join(manifest) + '\n')
print('Prepared', len(manifest) // 3, 'weight downloads')