#!/usr/bin/env python
"""Resolve-then-fetch sweep over every trainable source named in the four dataset docs.
Search HF by name -> pick best match -> download. NO-TRAIN / LICENCE-BLOCKED are
recorded but NEVER downloaded into the training tree."""
import os, json, time, sys
from concurrent.futures import ThreadPoolExecutor, as_completed
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
from huggingface_hub import HfApi, snapshot_download
API=HfApi(); TOK=os.environ.get("HF_TOKEN")
OUT="/workspace/aditi/webdata"; os.makedirs(OUT,exist_ok=True)

# (search term, local name, max_gb)  -- TRAINABLE only
WANT=[
 # Tier A backbones / CPT
 ("ai4bharat/sangraha","sangraha",40),("ai4bharat/IndicCorpV2","indiccorp_v2",40),
 ("ai4bharat/BPCC","bpcc",25),
 # Kashmir / Pakistan / low-resource
 ("ai4bharat/indicvoices_r","indicvoices_r",8),
 ("Ax-to-Grind Urdu fake news","ax_to_grind_urdu",2),
 ("urdu fake news bend the truth","urdufake_btt",2),
 ("SiNER sindhi ner","siner",2),("pashto offensive language","nlpashto",2),
 ("shahmukhi ner","shahmukhi_ner",2),
 # China tier
 ("CHEF chinese fact verification","chef_zh",3),("CFEVER chinese fever","cfever_zh",3),
 ("tibetan news classification","tncc_bo",2),("cmmlu","cmmlu",2),
 # Bangladesh tier
 ("Vashantor bangla dialect","vashantor",2),("BD-SHS bangla hate","bdshs",2),
 ("ANCHOLIK-NER bangla dialect","ancholik_ner",2),("BanFakeNews","banfakenews_v1",3),
 # OSINT task heads
 ("HiNER hindi ner","hiner",3),("MACD sharechat abuse","macd",2),
 ("SentiMix hinglish sentiment","sentimix",1),
 ("india hate speech superset","india_hate_superset",1),
 ("HASOC hindi offensive","hasoc_extra",2),
 ("CONSTRAINT hindi hostile posts","constraint_hostile",1),
 ("FACTIFY claim verification","factify2",5),
 ("Memotion meme sentiment","memotion3",4),
 ("code mixed offensive meme","cm_off_meme",2),
 ("L3Cube IndicNews","l3cube_indicnews",5),
 ("IndicXTREME","indicxtreme",5),
 ("Anudesh indic instruction","anudesh",2),
 ("IndicAlign toxic","indicalign_toxic",3),
 ("dakshina transliteration","dakshina",3),
 ("MILU indian language understanding","milu_EVALONLY",2),
]
# Recorded, deliberately NOT fetched into the training tree
FIREWALL={"FLORES-200":"NO-TRAIN (publisher)","Bhasha-Abhijnaanam":"NO-TRAIN (IndicLID bench)",
 "IndicGenBench":"NO-TRAIN (Tier-6 firewall)","IndicParam":"EVAL-ONLY (P4 regression tracker)",
 "IndicNLP Corpus":"EXCLUDED CC-BY-NC-SA","Balochi/Saraiki (BECO,NAWA-E-WATAN,Jhoke)":"LICENCE-BLOCKED CC-BY-NC",
 "L3Cube MahaSent/MahaHate/MahaNER":"eval-only / NC — owner approval needed",
 "WikiANN/PAN-X":"SILVER — warm-start only, never eval"}

def resolve(term):
    try:
        r=list(API.list_datasets(search=term, limit=5, sort="downloads", direction=-1))
        if not r: r=list(API.list_datasets(search=term.split("/")[-1], limit=5,
                                           sort="downloads", direction=-1))
        return [d.id for d in r]
    except Exception: return []

def grab(item):
    term,name,maxgb=item
    t0=time.time()
    cands=[term] if "/" in term else []
    cands+= [c for c in resolve(term) if c not in cands]
    for repo in cands[:3]:
        dest=os.path.join(OUT,name)
        try:
            info=API.dataset_info(repo, token=TOK)
            sz=sum((s.size or 0) for s in (info.siblings or []))
            snapshot_download(repo_id=repo, repo_type="dataset", local_dir=dest,
                              max_workers=6, token=TOK)
            b=sum(os.path.getsize(os.path.join(dp,f)) for dp,_,fs in os.walk(dest) for f in fs)
            if b> maxgb*1e9:
                return {"want":term,"repo":repo,"status":"oversize_kept","bytes":b,"secs":round(time.time()-t0)}
            return {"want":term,"repo":repo,"local":name,"status":"verified","bytes":b,
                    "secs":round(time.time()-t0)}
        except Exception as e:
            last=f"{type(e).__name__}: {str(e)[:140]}"
            continue
    return {"want":term,"status":"failed","tried":cands[:3],"error":locals().get("last","no candidate")}

if __name__=="__main__":
    res=[]
    with ThreadPoolExecutor(max_workers=6) as ex:
        futs={ex.submit(grab,w):w for w in WANT}
        for i,f in enumerate(as_completed(futs),1):
            r=f.result(); res.append(r)
            print(f"[{i}/{len(WANT)}] {r['status']:14s} {r.get('repo',r['want'])[:58]:58s} "
                  f"{r.get('bytes',0)/1e6:8.0f}MB {r.get('error','')[:70]}",flush=True)
    ok=[r for r in res if r['status'].startswith('verified') or r['status']=='oversize_kept']
    json.dump({"generated":time.strftime("%Y-%m-%dT%H:%M:%SZ",time.gmtime()),
               "verified":len(ok),"failed":len(res)-len(ok),
               "bytes":sum(r.get('bytes',0) for r in ok),
               "firewall_not_fetched":FIREWALL,"items":res},
              open("/workspace/aditi/ghdata/_reports/sweep_ledger.json","w"),indent=2)
    print(f"\nSWEEP verified={len(ok)} failed={len(res)-len(ok)} "
          f"{sum(r.get('bytes',0) for r in ok)/1e9:.1f}GB",flush=True)