"""Large public typed-decision corpora -> MM-Jev train records. Each builder returns a list of records; pipeline.py caches
every part to disk, so a crash / disconnect only redoes the part in flight.

Benchmark sources are removed by name (tasksource): AG News, emotion (DAIR), Banking77, MASSIVE, XNLI.
"""
import ast, json, random, re
from datasets import load_dataset
from mmjev import Seg, options_of

BANNED = re.compile(r"(^|/)(ag_news|emotion|banking77|massive|xnli)(/|$)|multilingual/massive|multilingual/xnli", re.I)
MAX_OPTS, MAX_STATE = 100, 3500


def _rec(task, state, qs, targets):
    ys = [int(max(range(len(t)), key=lambda i: t[i])) for t in targets]
    return dict(task=task, modality="text", split="train", state=[Seg("text", str(state)[:MAX_STATE])], qs=qs, ys=ys,
                targets=targets, raw=None)


def _norm(t):
    t = [max(0.0, float(x)) for x in t]
    s = sum(t)
    return [x / s for x in t] if s > 0 else None


def _parse(x):
    if isinstance(x, (list, dict)):
        return x
    try:
        return json.loads(x)
    except Exception:
        return ast.literal_eval(x)


def row_question(kind, question, options, target):
    """Row-schema (Open-Jev / SargeDev / tasksource): options list + target distribution -> (question, target)."""
    options = [str(o) for o in options]
    if kind == "noul" and len(target) == 1:
        target = [1 - float(target[0]), float(target[0])]
        options = ["false", "true"]
    target = _norm(target)
    if target is None or len(options) != len(target) or not (2 <= len(options) <= MAX_OPTS):
        return None
    if kind == "noul":
        p = target[-1] if len(target) == 2 else target[0]
        return {"type": "noul", "instructions": question}, [1 - p, p]
    if kind == "score":
        return {"type": "score", "instructions": question, "criteria": [o[:160] for o in options]}, target
    crit = {}
    for i, o in enumerate(options):
        if ": " in o and len(o.split(": ", 1)[0]) <= 40:
            k, v = o.split(": ", 1)
        elif len(o) > 60:
            k, v = f"option {i + 1}", o
        else:
            k, v = o, ""
        crit[k.strip()] = v[:300]
    if len(crit) != len(options):
        return None
    return {"type": "choice", "instructions": question, "criteria": crit}, target


def grouped_rows(rows, task, n_groups, rng, state_key="state_json"):
    groups = {}
    for r in rows:
        groups.setdefault(r["group_id"], []).append(r)
    keys = list(groups); rng.shuffle(keys)
    out = []
    for g in keys[:n_groups]:
        qs, ts, st = [], [], None
        for r in groups[g][:6]:
            st = st if st is not None else _parse(r[state_key]) if r[state_key] else ""
            got = row_question(r["kind"], r["question"], _parse(r["options"]), _parse(r["target"]))
            if got:
                qs.append(got[0]); ts.append(got[1])
        if qs:
            st = st if isinstance(st, str) else json.dumps(st, ensure_ascii=False)
            out.append(_rec(task, st, qs, ts))
    return out


def openjev_v11(n_ctrl=5000, n_wanli=1500, seed=11):
    rng = random.Random(seed)
    ds = load_dataset("ZefanCai/Open-Jev-v1.1", "community-hard-mix-v2-redistributable", split="train")
    rows = ds.to_list()
    wanli = [r for r in rows if "wanli" in (r["source"] or "").lower()]
    ctrl = [r for r in rows if "wanli" not in (r["source"] or "").lower()]
    return grouped_rows(ctrl, "openjev_v11", n_ctrl, rng) + grouped_rows(wanli, "openjev_v11_wanli", n_wanli, rng)


def openjev_v2(n=3000, seed=12):
    ds = load_dataset("ZefanCai/Open-Jev", "release-v2-redistributable", split="train")
    return grouped_rows(ds.to_list(), "openjev_v2", n, random.Random(seed))


def bev(cfgs=(("default", 4000), ("hard_50k", 3000), ("numeric_temporal", 1500), ("skills", 2000)), seed=13):
    out = []
    for cfg, n in cfgs:
        ds = load_dataset("avbiswas/bev-decision", cfg, split="train").shuffle(seed=seed).select(range(n))
        for ex in ds:
            qs_raw = _parse(ex["questions_json"])
            qs, ts = [], []
            for qid, q in list(qs_raw.items())[:6]:
                t, crit, lab = q["type"], q.get("criteria"), q.get("label")
                if t == "noul":
                    p = float(lab) if isinstance(lab, (int, float)) and not isinstance(lab, bool) else float(bool(lab))
                    qs.append({"type": "noul", "instructions": q["instructions"]}); ts.append([1 - p, p])
                elif t == "score":
                    crit = list(crit)
                    lv = int(lab) if lab is not None else 0
                    if not (0 <= lv < len(crit)):
                        continue
                    qs.append({"type": "score", "instructions": q["instructions"], "criteria": [str(c) for c in crit]})
                    ts.append([1.0 if i == lv else 0.0 for i in range(len(crit))])
                else:
                    crit = crit if isinstance(crit, dict) else {str(c): "" for c in crit}
                    keys = list(crit)
                    if str(lab) not in keys or not (2 <= len(keys) <= MAX_OPTS):
                        continue
                    qs.append({"type": "choice", "instructions": q["instructions"], "criteria": {k: str(v or "") for k, v in crit.items()}})
                    ts.append([1.0 if k == str(lab) else 0.0 for k in keys])
            if qs:
                out.append(_rec(f"bev:{cfg}", ex["state"], qs, ts))
    return out


def tasksource(n=7000, shards=(0, 4, 8), many_frac=0.45, seed=14):
    """Non-streaming shards (the streaming shuffle walks whole shards); benchmark sources dropped by name;
    rows with >= 10 options over-sampled so the K distribution covers intent-style label spaces."""
    rng = random.Random(seed)
    files = [f"hf://datasets/tasksource/tasksource-jev-typed-decisions/data/train-{i:05d}-of-00012.parquet" for i in shards]
    ds = load_dataset("parquet", data_files=files, split="train")
    ds = ds.filter(lambda b: [not BANNED.search(s or "") for s in b["source"]], batched=True)
    many, few = [], []
    for ex in ds.shuffle(seed=seed).select(range(min(len(ds), 60000))):
        try:
            opts = _parse(ex["options"]) if ex["options"] else ["false", "true"]
            tgt = _parse(ex["target"])
        except Exception:
            continue
        got = row_question(ex["kind"], ex["question"], opts, tgt if isinstance(tgt, list) else [1 - tgt, tgt])
        if not got:
            continue
        r = _rec(f"tasksource:{(ex['source'] or '?').split('/')[0]}", ex["state"] or "", [got[0]], [got[1]])
        (many if len(got[1]) >= 10 else few).append(r)
    k_many = min(len(many), int(n * many_frac))
    return rng.sample(many, k_many) + rng.sample(few, min(len(few), n - k_many))


def jev_decisions(n=2000, seed=15):
    """Agent tool / action choice: candidates -> options, target candidate -> gold."""
    rng = random.Random(seed)
    ds = load_dataset("samatv256/jev-decisions-v1", split="validation", streaming=True)
    out, seen = [], 0
    for ex in ds:
        seen += 1
        if seen > 60000 or len(out) >= n:
            break
        if ex["status"] != "trainable" or rng.random() > 0.25:
            continue
        try:
            cands = _parse(ex["candidates"]); tgt = _parse(ex["target"])
        except Exception:
            continue
        if not (2 <= len(cands) <= 40) or not isinstance(tgt, dict) or not tgt.get("candidate_id"):
            continue
        keys = [c.get("name") or c.get("id") for c in cands]
        ids = [c.get("id") for c in cands]
        if tgt["candidate_id"] not in ids or len(set(keys)) != len(keys):
            continue
        crit = {k: str(c.get("description") or "")[:200] for k, c in zip(keys, cands)}
        st = ex["state"] if isinstance(ex["state"], str) else json.dumps(ex["state"], ensure_ascii=False)
        st = st[-MAX_STATE:]                              # keep the most recent part of long agent traces
        g = ids.index(tgt["candidate_id"])
        out.append(_rec(f"jev_decisions:{ex['decision_type']}", st,
                        [{"type": "choice", "instructions": "Which action or tool should the agent use next?", "criteria": crit}],
                        [[1.0 if i == g else 0.0 for i in range(len(keys))]]))
    return out


def sargedev(n_v3=15000, n_v1=4000, seed=16):
    ds = load_dataset("json", data_files="hf://datasets/SargeDev/jev-distill-corpus-v3/train.jsonl", split="train")
    rng = random.Random(seed)
    by = {"yuri_v3": [], "yuri_v1": []}
    for i, s in enumerate(ds["source"]):
        if s in by:
            by[s].append(i)
    idx = rng.sample(by["yuri_v3"], min(n_v3, len(by["yuri_v3"]))) + rng.sample(by["yuri_v1"], min(n_v1, len(by["yuri_v1"])))
    out = []
    for ex in ds.select(idx):
        got = row_question(ex["kind"], ex["question"], ex["options"], ex["target"])
        if got:
            out.append(_rec(f"sargedev:{ex['source']}", ex["state"], [got[0]], [got[1]]))
    return out


def clinc_full(n=1200, seed=17):
    """CLINC150 with the FULL 150-intent label space (K matched to many-label intent routing)."""
    ds = load_dataset("clinc/clinc_oos", "plus", split="train").shuffle(seed=seed)
    names = ds.features["intent"].names
    labels = [x for x in names if x != "oos"]
    ds = ds.filter(lambda e: names[e["intent"]] != "oos").select(range(n))
    crit = {x.replace("_", " "): "" for x in labels}
    keys = list(crit)
    return [_rec("clinc150_full", ex["text"], [{"type": "choice", "instructions": "Which single label best describes the input text?",
                                                  "criteria": crit}],
                 [[1.0 if k == names[ex["intent"]].replace("_", " ") else 0.0 for k in keys]]) for ex in ds]


BUILDERS = {"openjev_v11": openjev_v11, "openjev_v2": openjev_v2, "bev": bev, "tasksource": tasksource,
            "jev_decisions": jev_decisions, "sargedev": sargedev, "clinc_full": clinc_full}