#!/usr/bin/env python3
"""Fetch pinned training files from the JSON index; dry-run unless --execute."""
import argparse,hashlib,json,shutil
from pathlib import Path,PurePosixPath


def safe(value):
 p=PurePosixPath(value)
 if p.is_absolute() or '..' in p.parts or not p.parts:raise ValueError('Unsafe relative data path')
 return Path(*p.parts)

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

def main():
 p=argparse.ArgumentParser(description=__doc__)
 directory=Path(__file__).absolute().parent
 default_index=directory/'TRAINING_DATA_INDEX.json'
 if not default_index.exists():default_index=directory.parent/'TRAINING_DATA_INDEX.json'
 p.add_argument('--index',type=Path,default=default_index)
 p.add_argument('--dataset',default='legacy_mix6_sft')
 p.add_argument('--output',type=Path,default=Path('training_data_download'))
 p.add_argument('--execute',action='store_true');a=p.parse_args()
 index=json.loads(a.index.read_text());datasets={d['dataset_id']:d for d in index['datasets']}
 if a.dataset=='all_private':ids={r['artifact_id'] for r in index['private_artifacts']}
 elif a.dataset in datasets:ids=set(datasets[a.dataset].get('artifact_ids',[]))
 else:raise ValueError('Unknown dataset ID; inspect index.datasets')
 rows=[r for r in index['private_artifacts'] if r['artifact_id'] in ids]
 if not rows:raise ValueError('No private snapshot artifacts for this dataset; inspect public download location')
 source=index['private_snapshot']
 if not source['repo_id'].startswith('diffusion-gif/training-data-vault-') or source['repo_type']!='dataset' or not source['private']:raise ValueError('Unexpected private data provider')
 print(json.dumps({'dataset':a.dataset,'files':len(rows),'bytes':sum(r['bytes'] for r in rows),'source_repo':source['repo_id'],'revision':source['revision'],'login_required':True,'execute':a.execute},indent=2))
 if not a.execute:return
 from huggingface_hub import get_token,hf_hub_download
 token=get_token()
 if not token:raise RuntimeError('Run hf auth login locally using an account authorized for the private vault. Never put tokens in this file or chat.')
 a.output.mkdir(parents=True,exist_ok=True)
 for r in rows:
  target=a.output/safe(r['path'])
  if target.exists():
   if not target.is_file() or target.stat().st_size!=r['bytes'] or sha(target)!=r['sha256']:raise ValueError('Existing user file differs; will not overwrite')
   continue
  cached=Path(hf_hub_download(repo_id=source['repo_id'],repo_type='dataset',revision=source['revision'],filename=r['path'],token=token,endpoint='https://huggingface.co'))
  if cached.stat().st_size!=r['bytes'] or sha(cached)!=r['sha256']:raise ValueError('Downloaded file SHA256 mismatch')
  target.parent.mkdir(parents=True,exist_ok=True);temporary=target.with_name(target.name+'.copy-partial')
  shutil.copyfile(cached,temporary)
  if sha(temporary)!=r['sha256']:raise ValueError('Copied file SHA256 mismatch')
  temporary.replace(target);print('VERIFIED',r['path'],flush=True)
 print('Done: all selected frozen files downloaded and SHA256 verified. See restore_legacy_mix6.py for a relocatable SFT working copy.')
if __name__=='__main__':main()