"""Download the pinned base model, reference model, training data and benchmark; record hashes."""
import concurrent.futures
from huggingface_hub import snapshot_download
from common import *

SPECS = {
    'base': (*BASE, ['config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json', 'special_tokens_map.json', 'README.md']),
    'reference': (*REFERENCE, ['config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json', 'README.md',
                               'eval_results.json', 'benchmark_predictions.npz', 'eval/tool_probes.json', 'eval/fresh_probes.json']),
    'data': (*DATASET, ['*.parquet', 'README.md']),
    'benchmark': (*BENCHMARK, ['*.parquet', 'README.md']),
}

def get(item):
    name, (repo, kind, rev, patterns) = item
    path = snapshot_download(repo, repo_type=kind, revision=rev, local_dir=DATA/'raw'/name, allow_patterns=patterns, max_workers=8)
    event('downloaded', name=name, repo=repo, revision=rev)
    return name, {'repo': repo, 'kind': kind, 'revision': rev, 'path': path}

def main():
    if (DATA/'sources.json').exists(): print('sources ready'); return
    assert (TEACHER_LOGITS/'teacher_logits.npy').exists(), 'external SSD archive must be mounted read-only'
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
        manifest = dict(pool.map(get, SPECS.items()))
    for name in ['base', 'reference']:
        manifest[name]['weights_sha256'] = sha(Path(manifest[name]['path'])/'model.safetensors')
    for name, f in [('data', 'train.parquet'), ('data', 'validation.parquet'), ('benchmark', 'test.parquet')]:
        manifest[name].setdefault('files', {})[f] = sha(Path(manifest[name]['path'])/f)
    manifest['teacher'] = {'name': 'Private Auto 3B classifier (SmolLM3-3B-Base, archived, unpublished)', 'weights_sha256': TEACHER_SHA,
                           'cached_train_logits': str(TEACHER_LOGITS/'teacher_logits.npy'),
                           'cached_aug_logits': str(AUG/'aug_teacher_logits.npy')}
    atomic_json(DATA/'sources.json', manifest)
    event('sources_ready', **{k: v.get('revision') for k, v in manifest.items()})

if __name__ == '__main__': main()