#!/usr/bin/env python
"""build_corpus.py - bucket stages -> aditi_v8.jsonl. Every payload validated at emit
time against the gate's own schema (contract.Emitter). Budgets in SUPERVISED tokens."""
import os, sys, json, glob, zipfile, io, random, argparse, re
sys.path.insert(0,"/workspace/aditi")
import contract as C
import pandas as pd
GH="/workspace/aditi/ghdata"
random.seed(1947)
def clip(t,n=1500):
    t=(t or "").strip(); return t[:n]
def pq(p):
    try: return pd.read_parquet(p)
    except Exception: return None

# supervised-token budgets (Emitter charges assistant string only)
BUDGET={"faithful_gen":700_000,"schema_lock":420_000,"halluc":900_000,"claim":320_000,
        "halluc_indic":520_000,"ner":300_000,"general":180_000}
CAPS  ={"faithful_gen":12000,"schema_lock":6310,"halluc":26000,"claim":9000,
        "halluc_indic":14000,"ner":9000,"general":16000}
count={}
def room(b): return count.get(b,0) < CAPS[b]
def took(b): count[b]=count.get(b,0)+1

# ---------------------------------------------------------------- faithful_gen (G2)
def do_faithful_gen(E):
    d=pq(f"{GH}/RAGTruth-processed/data/train-00000-of-00001.parquet")
    if d is not None:
        def _hl(x):
            if x is None: return []
            if isinstance(x,str):
                try: return json.loads(x)
                except Exception: return []
            return list(x)
        d=d.copy(); d["_hl"]=d["hallucination_labels"].apply(_hl)
        sub=d[d["_hl"].apply(len)==0]
        for _,r in sub.iterrows():
            if not room("faithful_gen"): break
            u=C.P_GROUND_GEN.format(iid=C.content_hash(str(r["context"]))[:12],
                source=clip(str(r["context"]),1500), question=clip(str(r["query"]),400))
            if E.emit("faithful_gen","ragtruth-grounded",u,{"answer":clip(str(r["output"]),1200)},
                      C.GROUNDED_ANSWER_SCHEMA, heldout_text=str(r["context"])): took("faithful_gen")
    a=pq(glob.glob(f"{GH}/asqa/data/train-*.parquet")[0]) if glob.glob(f"{GH}/asqa/data/train-*.parquet") else None
    if a is not None:
        for _,r in a.iterrows():
            if not room("faithful_gen"): break
            qas=r.get("qa_pairs"); ann=r.get("annotations")
            ctx=" ".join([str(q.get("context","")) for q in (qas if qas is not None else [])])[:1500]
            ans=""
            try: ans=str(ann[0]["long_answer"])
            except Exception: ans=str(r.get("sample_id",""))
            if not ctx or not ans: continue
            u=C.P_GROUND_GEN.format(iid=C.content_hash(ctx)[:12],source=ctx,question=clip(str(r["ambiguous_question"]),400))
            if E.emit("faithful_gen","asqa",u,{"answer":clip(ans,1200)},C.GROUNDED_ANSWER_SCHEMA,
                      heldout_text=ctx): took("faithful_gen")
    for lg in ("ta","bn","gu","hi","kn","te","mr"):
        f=f"{GH}/Indic-Rag-Suite/data/{lg}_data.parquet"
        if not os.path.exists(f): continue
        d=pq(f)
        if d is None: continue
        for _,r in d.head(4000).iterrows():
            if not room("faithful_gen"): break
            ctx=clip(str(r.get("paragraph","")),1500); q=clip(str(r.get("question","")),400)
            ans=str(r.get("answer",""))
            if len(ctx)<40 or not q or not ans: continue
            u=C.P_GROUND_GEN.format(iid=C.content_hash(ctx)[:12],source=ctx,question=q)
            if E.emit("faithful_gen",f"indic-rag-{lg}",u,{"answer":clip(ans,900)},
                      C.GROUNDED_ANSWER_SCHEMA,heldout_text=ctx): took("faithful_gen")

# ---------------------------------------------------------------- halluc (G1 binding)
def do_halluc(E):
    d=pq(f"{GH}/RAGTruth-processed/data/train-00000-of-00001.parquet")
    if d is not None:
        def _hl(x):
            if x is None: return []
            if isinstance(x,str):
                try: return json.loads(x)
                except Exception: return []
            return list(x)
        d=d.copy(); d["_hl"]=d["hallucination_labels"].apply(_hl)
        for _,r in d.iterrows():
            if not room("halluc"): break
            labs=r["_hl"]; g=(len(labs)==0)
            spans=[] if g else [str(x.get("text",""))[:200] for x in labs if isinstance(x,dict)][:6]
            u=C.P_HALLUC.format(iid=C.content_hash(str(r["context"]))[:12],
                source=clip(str(r["context"]),1500),response=clip(str(r["output"]),900))
            if E.emit("halluc","ragtruth-train",u,{"grounded":bool(g),"unsupported_spans":spans},
                      C.HALLUC_SCHEMA,heldout_text=str(r["context"])): took("halluc")
    d=pq(f"{GH}/LLM-AggreFact_sub/data/train-00000-of-00001.parquet")
    if d is not None:
        for _,r in d.iterrows():
            if not room("halluc"): break
            g=bool(r["label"])
            u=C.P_HALLUC.format(iid=C.content_hash(str(r["doc"]))[:12],
                source=clip(str(r["doc"]),1500),response=clip(str(r["claim"]),600))
            if E.emit("halluc","LLM-AggreFact_sub",u,{"grounded":g,"unsupported_spans":[] if g else [clip(str(r["claim"]),200)]},
                      C.HALLUC_SCHEMA,heldout_text=str(r["doc"])): took("halluc")
    for task,src,kn,qn in (("qa","halueval-qa","knowledge","question"),
                           ("dialogue","halueval-dialogue","knowledge","dialogue_history"),
                           ("summarization","halueval-summarization","document","document")):
        fs=glob.glob(f"{GH}/HaluEval/{task}/data-*.parquet")
        if not fs: continue
        d=pq(fs[0])
        if d is None: continue
        cols=set(d.columns)
        rk=[c for c in ("right_answer","right_response","right_summary") if c in cols]
        hk=[c for c in ("hallucinated_answer","hallucinated_response","hallucinated_summary") if c in cols]
        kk=kn if kn in cols else ("document" if "document" in cols else None)
        if not rk or not hk or not kk: continue
        for _,r in d.head(4000).iterrows():
            for key,g in ((rk[0],True),(hk[0],False)):
                if not room("halluc"): break
                src_t=clip(str(r[kk]),1500)
                u=C.P_HALLUC.format(iid=C.content_hash(src_t+key)[:12],source=src_t,
                                    response=clip(str(r[key]),700))
                if E.emit("halluc",src,u,{"grounded":g,"unsupported_spans":[] if g else [clip(str(r[key]),200)]},
                          C.HALLUC_SCHEMA,heldout_text=src_t): took("halluc")

def do_halluc_indic(E):
    for lg in ("ta","ur","kn","te","ml","as","gu","bn","hi","mr","pa","or"):
        f=f"{GH}/Indic-Rag-Suite/data/{lg}_data.parquet"
        if not os.path.exists(f): continue
        d=pq(f)
        if d is None: continue
        rows=d.head(3000)
        for i,(_,r) in enumerate(rows.iterrows()):
            if not room("halluc_indic"): break
            ctx=clip(str(r.get("paragraph","")),1200); ans=str(r.get("answer",""))
            if len(ctx)<40 or not ans: continue
            if i%2==0:
                g=True; resp=ans
            else:
                g=False; resp=str(rows.iloc[(i+7)%len(rows)].get("answer",""))
                if not resp or resp==ans: continue
            u=C.P_HALLUC.format(iid=C.content_hash(ctx+str(i))[:12],source=ctx,response=clip(resp,600))
            if E.emit("halluc_indic",f"indic-halluc-{lg}",u,
                      {"grounded":g,"unsupported_spans":[] if g else [clip(resp,200)]},
                      C.HALLUC_SCHEMA,heldout_text=ctx): took("halluc_indic")

# ---------------------------------------------------------------- claim (G1)
def do_claim(E):
    d=pq(f"{GH}/averitec/data/train-00000-of-00001.parquet")
    if d is not None:
        for _,r in d.iterrows():
            if not room("claim"): break
            v=C.AVERITEC_MAP.get(str(r["label"]))
            if not v: continue
            ev=""
            try: ev=" ".join([f"{q.get('question','')} {q.get('answers','')}" for q in r["questions"]])[:1500]
            except Exception: ev=str(r.get("justification",""))[:1500]
            if not ev.strip(): continue
            u=C.P_CLAIM.format(iid=C.content_hash(str(r["claim"]))[:12],claim=clip(str(r["claim"]),500),evidence=clip(ev,1500))
            if E.emit("claim","averitec-train",u,{"verdict":v},C.CLAIM_SCHEMA,heldout_text=str(r["claim"])): took("claim")
    for f in sorted(glob.glob(f"{GH}/x-fact/all_languages/train-*.csv")):
        try: d=pd.read_csv(f)
        except Exception: continue
        for _,r in d.iterrows():
            if not room("claim"): break
            v=C.canon("verdict",str(r.get("label","")),strict=False)
            if not v: continue
            ev=" ".join([str(r.get(f"evidence_{i}","")) for i in range(1,6)])[:1500]
            if len(ev.strip())<20: continue
            u=C.P_CLAIM.format(iid=C.content_hash(str(r["claim"]))[:12],claim=clip(str(r["claim"]),500),evidence=clip(ev,1500))
            if E.emit("claim","x-fact",u,{"verdict":v},C.CLAIM_SCHEMA,heldout_text=str(r["claim"])): took("claim")
    d=pq(f"{GH}/climate_fever/data/test-00000-of-00001.parquet") if glob.glob(f"{GH}/climate_fever/data/*.parquet") else None
    if d is None:
        g=glob.glob(f"{GH}/climate_fever/**/*.parquet",recursive=True); d=pq(g[0]) if g else None
    if d is not None:
        CF={0:"supported",1:"refuted",2:"not_enough_evidence",3:"conflicting_evidence"}
        for _,r in d.iterrows():
            if not room("claim"): break
            v=CF.get(int(r["claim_label"])) if "claim_label" in d.columns else None
            if not v: continue
            try: ev=" ".join([e["evidence"] for e in r["evidences"]])[:1500]
            except Exception: continue
            u=C.P_CLAIM.format(iid=C.content_hash(str(r["claim"]))[:12],claim=clip(str(r["claim"]),500),evidence=clip(ev,1500))
            if E.emit("claim","climate-fever",u,{"verdict":v},C.CLAIM_SCHEMA,heldout_text=str(r["claim"])): took("claim")
    g=glob.glob(f"{GH}/pubhealth-converted/**/*.parquet",recursive=True)
    if g:
        d=pq(g[0])
        if d is not None:
            for _,r in d.iterrows():
                if not room("claim"): break
                v=C.canon("verdict",str(r.get("label","")),strict=False)
                ev=str(r.get("main_text",r.get("evidence","")))[:1500]
                if not v or len(ev.strip())<20: continue
                u=C.P_CLAIM.format(iid=C.content_hash(str(r.get("claim","")))[:12],
                                   claim=clip(str(r.get("claim","")),500),evidence=clip(ev,1500))
                if E.emit("claim","pubhealth",u,{"verdict":v},C.CLAIM_SCHEMA,heldout_text=str(r.get("claim",""))): took("claim")

# ---------------------------------------------------------------- general (G4)
def do_general(E):
    LET="ABCDEFGHIJ"
    for sub in sorted(os.listdir(f"{GH}/mmlu")):
        p=f"{GH}/mmlu/{sub}"
        if not os.path.isdir(p): continue
        fs=glob.glob(f"{p}/dev-*.parquet")+glob.glob(f"{p}/validation-*.parquet")
        for f in fs:
            d=pq(f)
            if d is None: continue
            for _,r in d.iterrows():
                if not room("general"): break
                ch=list(r["choices"]); ans=int(r["answer"])
                q=str(r["question"]); cs="\n".join(f"{LET[i]}. {c}" for i,c in enumerate(ch))
                u=C.P_MCQ.format(iid=C.content_hash(q)[:12],question=clip(q,700),choices=cs)
                if E.emit("general","mmlu-all",u,{"answer":LET[ans]},C.MCQ_SCHEMA,heldout_text=q): took("general")
    for nm,sub in (("ai2_arc-ARC-Easy","ARC-Easy"),("ai2_arc-ARC-Challenge","ARC-Challenge")):
        for f in glob.glob(f"{GH}/ai2_arc/{sub}/train-*.parquet")+glob.glob(f"{GH}/ai2_arc/{sub}/validation-*.parquet"):
            d=pq(f)
            if d is None: continue
            for _,r in d.iterrows():
                if not room("general"): break
                ch=r["choices"]; texts=list(ch["text"]); labs=list(ch["label"])
                key=str(r["answerKey"])
                if key not in labs: continue
                idx=labs.index(key)
                if idx>=10: continue
                cs="\n".join(f"{LET[i]}. {t}" for i,t in enumerate(texts))
                u=C.P_MCQ.format(iid=C.content_hash(str(r["question"]))[:12],question=clip(str(r["question"]),700),choices=cs)
                if E.emit("general",nm,u,{"answer":LET[idx]},C.MCQ_SCHEMA,heldout_text=str(r["question"])): took("general")
    for f in glob.glob(f"{GH}/mmlu-indic/**/*.parquet",recursive=True):
        lg=os.path.basename(f).split("-")[0][:2]
        d=pq(f)
        if d is None or "choices" not in d.columns: continue
        for _,r in d.iterrows():
            if not room("general"): break
            try:
                ch=list(r["choices"]); ans=int(r["answer"])
            except Exception: continue
            cs="\n".join(f"{LET[i]}. {c}" for i,c in enumerate(ch))
            u=C.P_MCQ.format(iid=C.content_hash(str(r["question"]))[:12],question=clip(str(r["question"]),700),choices=cs)
            if E.emit("general",f"mmlu-indic-{lg}",u,{"answer":LET[ans]},C.MCQ_SCHEMA,heldout_text=str(r["question"])): took("general")

STAGES={"faithful_gen":do_faithful_gen,"halluc":do_halluc,"halluc_indic":do_halluc_indic,
        "claim":do_claim,"general":do_general}
if __name__=="__main__":
    ap=argparse.ArgumentParser(); ap.add_argument("--out",required=True)
    ap.add_argument("--only",default=""); a=ap.parse_args()
    E=C.Emitter(a.out,BUDGET,tokenizer_dir="/workspace/models/gemma4-qat-unquantized")
    only=[s for s in a.only.split(",") if s] or list(STAGES)
    for s in only:
        print(f"--- stage {s}",flush=True)
        try: STAGES[s](E)
        except Exception as ex:
            import traceback; traceback.print_exc()
        print(f"    {s}: {count.get(s,0)} recs, {E.spent.get(s,0)} sup-tok",flush=True)
    E.close()
    rep=E.report(a.out+".report.txt"); print(rep)