#!/usr/bin/env python
"""SERIAL harvester: fetch -> extract to compact jsonl.gz -> DELETE raw. Disk stays flat.
One source at a time. Never runs concurrently with itself."""
import os,sys,json,gzip,glob,time,shutil,subprocess,zipfile,io,csv,re
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
from huggingface_hub import snapshot_download
TOK=os.environ.get("HF_TOKEN")
STAGE="/workspace/aditi/_stage"; HARV="/workspace/aditi/harvest"
LEDGER="/workspace/aditi/ghdata/_reports/harvest_ledger.json"
os.makedirs(HARV,exist_ok=True)
MAX_ROWS=int(os.environ.get("MAX_ROWS","1500000"))
MIN_FREE_GB=25
TEXTCOLS=re.compile(r"(text|sentence|content|body|claim|question|answer|comment|title|"
                    r"headline|tweet|post|review|source|context|doc|output|response|"
                    r"summary|prompt|instruction|input|target|translation|word|token|caption)",re.I)

def free_gb():
    s=os.statvfs("/"); return s.f_bavail*s.f_frsize/1e9

def rows_from_any(path):
    """Yield dicts from parquet/csv/tsv/json/jsonl/txt inside a directory tree."""
    import pandas as pd
    exts=("**/*.parquet","**/*.csv","**/*.tsv","**/*.jsonl","**/*.json","**/*.txt")
    files=[]
    for e in exts: files+=glob.glob(os.path.join(path,e),recursive=True)
    files=[f for f in files if ".cache" not in f and os.path.getsize(f)>0]
    files.sort(key=os.path.getsize)
    for f in files[:2000]:
        try:
            low=f.lower()
            if low.endswith(".parquet"):
                d=pd.read_parquet(f)
                for r in d.head(MAX_ROWS).to_dict("records"): yield r
            elif low.endswith((".csv",".tsv")):
                sep="\t" if low.endswith(".tsv") else ","
                for chunk in pd.read_csv(f,sep=sep,encoding="utf-8-sig",on_bad_lines="skip",
                                         chunksize=50000,dtype=str,engine="python"):
                    for r in chunk.to_dict("records"): yield r
            elif low.endswith(".jsonl"):
                with open(f,encoding="utf-8",errors="replace") as fh:
                    for i,l in enumerate(fh):
                        if i>=MAX_ROWS: break
                        try: yield json.loads(l)
                        except Exception: pass
            elif low.endswith(".json"):
                with open(f,encoding="utf-8",errors="replace") as fh:
                    try: d=json.load(fh)
                    except Exception: continue
                if isinstance(d,list):
                    for r in d[:MAX_ROWS]:
                        if isinstance(r,dict): yield r
            elif low.endswith(".txt") and os.path.getsize(f)<400_000_000:
                with open(f,encoding="utf-8",errors="replace") as fh:
                    for i,l in enumerate(fh):
                        if i>=MAX_ROWS: break
                        l=l.strip()
                        if len(l)>20: yield {"text":l}
        except Exception:
            continue

def compact(r):
    """Keep only text-bearing fields + short labels. Drops blobs/images."""
    o={}
    for k,v in (r.items() if isinstance(r,dict) else []):
        if v is None: continue
        if isinstance(v,(bytes,bytearray)): continue
        s=v if isinstance(v,str) else (json.dumps(v,ensure_ascii=False,default=str)
                                       if isinstance(v,(list,dict)) else str(v))
        if len(s)>6000: s=s[:6000]
        if TEXTCOLS.search(str(k)) or len(s)<=64:
            o[str(k)]=s
    return o if any(len(v)>15 for v in o.values()) else None

def extract(name, srcdir, meta):
    out=os.path.join(HARV,f"{name}.jsonl.gz"); n=0
    with gzip.open(out,"wt",encoding="utf-8") as w:
        w.write(json.dumps({"_meta":meta})+"\n")
        for r in rows_from_any(srcdir):
            c=compact(r)
            if not c: continue
            w.write(json.dumps(c,ensure_ascii=False)+"\n"); n+=1
            if n>=MAX_ROWS: break
    return out,n,os.path.getsize(out)

def do_hf(repo,name,pats=None,lic=""):
    d=os.path.join(STAGE,name)
    snapshot_download(repo_id=repo,repo_type="dataset",local_dir=d,
                      allow_patterns=pats,max_workers=8,token=TOK)
    return d,{"source":repo,"kind":"hf","licence":lic}
def do_git(url,name,lic=""):
    d=os.path.join(STAGE,name)
    subprocess.run(["git","clone","--depth","1","--quiet",url,d],
                   capture_output=True,timeout=1200,check=True)
    return d,{"source":url,"kind":"git","licence":lic}

def run(queue):
    led=[]
    if os.path.exists(LEDGER):
        try: led=json.load(open(LEDGER))["items"]
        except Exception: led=[]
    done={i["name"] for i in led if i.get("status")=="harvested"}
    for kind,ref,name,pats,lic in queue:
        if name in done:
            print(f"skip {name} (already harvested)",flush=True); continue
        if free_gb()<MIN_FREE_GB:
            print(f"ABORT: only {free_gb():.0f}G free",flush=True); break
        t0=time.time(); shutil.rmtree(STAGE,ignore_errors=True); os.makedirs(STAGE,exist_ok=True)
        try:
            d,meta=(do_hf(ref,name,pats,lic) if kind=="hf" else do_git(ref,name,lic))
            raw=sum(os.path.getsize(os.path.join(dp,f)) for dp,_,fs in os.walk(d) for f in fs)
            out,n,sz=extract(name,d,meta)
            shutil.rmtree(STAGE,ignore_errors=True)              # <-- FLUSH
            led.append({"name":name,"ref":ref,"kind":kind,"licence":lic,"status":"harvested",
                        "rows":n,"raw_mb":round(raw/1e6),"kept_mb":round(sz/1e6,1),
                        "secs":round(time.time()-t0)})
            print(f"OK   {name:26s} rows={n:>7} raw={raw/1e6:7.0f}MB kept={sz/1e6:6.1f}MB "
                  f"free={free_gb():.0f}G",flush=True)
        except Exception as e:
            shutil.rmtree(STAGE,ignore_errors=True)
            led.append({"name":name,"ref":ref,"kind":kind,"status":"failed",
                        "error":f"{type(e).__name__}: {str(e)[:160]}"})
            print(f"FAIL {name:26s} {type(e).__name__}: {str(e)[:110]}",flush=True)
        json.dump({"generated":time.strftime("%Y-%m-%dT%H:%M:%SZ",time.gmtime()),
                   "items":led},open(LEDGER,"w"),indent=2)
    ok=[i for i in led if i.get("status")=="harvested"]
    print(f"\nHARVEST harvested={len(ok)} rows={sum(i.get('rows',0) for i in ok)} "
          f"kept={sum(i.get('kept_mb',0) for i in ok)/1024:.2f}GB free={free_gb():.0f}G",flush=True)
if __name__=="__main__":
    from harvest_queue import QUEUE
    run(QUEUE)