#!/usr/bin/env python
"""Reproduce the corpus from sources.json. Datasets are referenced by exact commit SHA,
never redistributed. Fetch -> extract -> flush, so disk stays flat."""
import os,sys,json,argparse,subprocess,shutil
def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--manifest",default="sources.json")
    ap.add_argument("--out",default="/workspace/aditi/media")
    ap.add_argument("--kind",choices=["media","text","all"],default="all")
    ap.add_argument("--dry-run",action="store_true")
    A=ap.parse_args()
    m=json.load(open(A.manifest)); S=m["sources"]
    if A.kind!="all": S=[s for s in S if s.get("kind")==A.kind]
    print(f"{len(S)} sources ({sum(1 for s in S if s.get('revision'))} SHA-pinned)")
    if A.dry_run:
        for s in S[:20]:
            print(f"  {s['host']:<12} {s['id']:<52} @{(s.get('revision') or 'HEAD')[:10]}")
        return
    os.environ.setdefault("HF_HOME","/workspace/.hf_home")
    from huggingface_hub import snapshot_download
    tok=os.environ.get("HF_TOKEN")
    for i,s in enumerate(S,1):
        dest=os.path.join(A.out,"_refetch",s.get("local") or s["id"].split("/")[-1])
        try:
            if s["host"]=="huggingface":
                snapshot_download(repo_id=s["id"],repo_type=s.get("repo_type","dataset"),
                                  revision=s.get("revision"),local_dir=dest,
                                  max_workers=8,token=tok)
            else:
                if os.path.isdir(dest): shutil.rmtree(dest,ignore_errors=True)
                subprocess.run(["git","clone","--depth","1","--quiet",s["id"],dest],
                               check=True,timeout=1800)
            print(f"[{i}/{len(S)}] OK {s['id']}",flush=True)
        except Exception as e:
            print(f"[{i}/{len(S)}] FAIL {s['id']}: {type(e).__name__} {str(e)[:100]}",flush=True)
    print("\nNow run: python code/harvest_media.py   (resize/resample + flush)")
    print("Then:    python build_corpus_mm.py --out corpus/aditi_mm_big.jsonl")
if __name__=="__main__": main()