"""Build data/{train,val,test,heldout}.jsonl + data/STATS.md from public HF datasets.

    python -m jevlike.data.build_dataset                       # full build (seed 13)
    python -m jevlike.data.build_dataset --seed 7 --max-per-source 8000
    python -m jevlike.data.build_dataset --sources ag_news,boolq --out /tmp/x   # quick subset

Pipeline: per-source converters (jevlike/data/sources_*.py) -> Records (validated, states
cleaned + truncated to 2000 chars) -> dedupe -> type-mix balancing (water-filling, so the
largest sources are trimmed first) -> split by example key (all questions about one state
land in the same split) with val/test quotas proportional to each source -> leak removal
(no state shared between splits; heldout states also checked by 8-word shingles) -> JSONL.

Deterministic given --seed: every random decision uses random.Random(f"{seed}|...").
Only writes data/{train,val,test,heldout}.jsonl and data/STATS.md (never data/synth, data/bench).
"""
from __future__ import annotations

import argparse
import hashlib
import os
import random
import re
import sys
import time
import traceback
from dataclasses import dataclass

from jevlike.data import sources_choice, sources_heldout, sources_noul, sources_score
from jevlike.data.common import clean, norm_key, truncate
from jevlike.data.sources_choice import Ctx
from jevlike.types import QTYPES, Record, write_jsonl

IN_DOMAIN = {**sources_choice.SOURCES, **sources_score.SOURCES, **sources_noul.SOURCES}
HELDOUT = dict(sources_heldout.SOURCES)
PRIMARY = {**{k: "choice" for k in sources_choice.SOURCES}, **{k: "score" for k in sources_score.SOURCES},
           **{k: "noul" for k in sources_noul.SOURCES}, "trec": "choice", "pavlick_formality": "score",
           "enron_spam": "noul"}
SPLIT_RANK = {"train": 0, "val": 1, "test": 2, "heldout": 3}
_UNFILLED = re.compile(r"\{(x|thing|target)\}")


@dataclass
class Item:
    rec: Record
    key: str        # example key (split unit)
    skey: str       # normalized state
    split: str = "train"


def log(*a):
    print(*a, file=sys.stderr, flush=True)


# ---------------------------------------------------------------- generation

def run_source(name: str, fn, seed: int, budget: int) -> list[Item]:
    ctx = Ctx(random.Random(f"{seed}|{name}"), name)
    items: list[Item] = []
    n_ex = 0
    for ex in fn(ctx):
        state = truncate(clean(ex.state))
        if not state:
            continue
        key = ex.key or norm_key(state)
        skey = norm_key(state)
        for j, (q, label) in enumerate(ex.qs):
            q.instructions = q.instructions.strip()
            assert q.instructions, (name, state[:80])
            assert not _UNFILLED.search(q.instructions), (name, q.instructions)
            rec = Record(f"{name}-{n_ex:06d}-{j}", name, state, q, int(label))
            rec.validate()
            items.append(Item(rec, key, skey))
        n_ex += 1
        if len(items) >= budget:
            break
    return items


def dedupe(items: list[Item]) -> list[Item]:
    """Drop exact duplicate questions about the same state; drop both copies if the
    duplicates disagree on the answer."""
    seen: dict[tuple, int] = {}
    answer: dict[tuple, str] = {}
    bad: set = set()
    keep: list[Item] = []
    for it in items:
        q = it.rec.question
        sig = (it.skey, q.type, q.instructions.lower(), tuple((o.key, o.text) for o in q.options))
        ans = q.options[it.rec.label].key if q.options else str(it.rec.label)
        if sig in seen:
            if answer[sig] != ans:
                bad.add(sig)
            continue
        seen[sig] = len(keep)
        answer[sig] = ans
        keep.append(it)
    return [it for it in keep if (it.skey, it.rec.question.type, it.rec.question.instructions.lower(),
                                  tuple((o.key, o.text) for o in it.rec.question.options)) not in bad]


def balance_types(items: list[Item], mix: dict[str, float], total: int, seed: int,
                  native_weight: float = 3.0) -> list[Item]:
    """Trim to `total` records with type shares `mix`. Within a type, every source is capped
    at the same level (water-filling, so the largest sources are trimmed first), except that
    sources whose native task is that type get a `native_weight`x higher cap (e.g. boolq /
    NLI statements outweigh noul statements derived from classification labels)."""
    by_t: dict[str, list[Item]] = {t: [] for t in QTYPES}
    for it in items:
        by_t[it.rec.question.type].append(it)
    feasible = min(len(by_t[t]) / mix[t] for t in QTYPES if mix[t] > 0)
    grand = min(total, int(feasible))
    out: list[Item] = []
    for t in QTYPES:
        target = int(round(mix[t] * grand))
        per_src: dict[str, list[Item]] = {}
        for it in by_t[t]:
            per_src.setdefault(it.rec.source, []).append(it)
        srcs = sorted(per_src)
        weights = [native_weight if PRIMARY.get(s) == t else 1.0 for s in srcs]
        level = _water_level([len(per_src[s]) for s in srcs], weights, target)
        for src, wgt in zip(srcs, weights):
            v = per_src[src]
            cap = int(level * wgt)
            if len(v) > cap:
                rng = random.Random(f"{seed}|balance|{src}|{t}")
                idx = sorted(rng.sample(range(len(v)), cap))
                v = [v[i] for i in idx]
            out.extend(v)
    return out


def _water_level(sizes: list[int], weights: list[float], target: int) -> float:
    """Largest level c with sum(min(size_i, int(c * w_i))) <= target."""
    def used(c):
        return sum(min(n, int(c * w)) for n, w in zip(sizes, weights))
    hi = max((n / w for n, w in zip(sizes, weights)), default=0.0)
    if used(hi) <= target:
        return hi
    lo = 0.0
    for _ in range(60):
        mid = (lo + hi) / 2
        if used(mid) <= target:
            lo = mid
        else:
            hi = mid
    return lo


# ---------------------------------------------------------------- splitting / leaks

def assign_splits(items: list[Item], val_size: int, test_size: int, seed: int) -> None:
    by_src: dict[str, list[Item]] = {}
    for it in items:
        by_src.setdefault(it.rec.source, []).append(it)
    n_total = len(items)
    for src in sorted(by_src):
        its = by_src[src]
        groups: dict[str, list[Item]] = {}
        for it in its:
            groups.setdefault(it.key, []).append(it)
        keys = list(groups)
        random.Random(f"{seed}|split|{src}").shuffle(keys)
        q_test = round(test_size * len(its) / n_total)
        q_val = round(val_size * len(its) / n_total)
        n_test = n_val = 0
        for k in keys:
            if n_test < q_test:
                split = "test"
                n_test += len(groups[k])
            elif n_val < q_val:
                split = "val"
                n_val += len(groups[k])
            else:
                split = "train"
            for it in groups[k]:
                it.split = split


def remove_leaks(items: list[Item], heldout: list[Item], shingle_n: int = 8) -> tuple[list[Item], dict]:
    """No example key / normalized state may occur in two splits: keep it in the
    highest-ranked split (heldout > test > val > train) and drop the other copies.
    In-domain states that share an 8-word shingle with a heldout state are dropped too."""
    rank: dict[str, int] = {}
    for it in items + heldout:
        r = SPLIT_RANK[it.split]
        for k in (it.key, it.skey):
            rank[k] = max(rank.get(k, 0), r)
    held_sh: set = set()
    for it in heldout:
        held_sh.update(_shingles(it.skey, shingle_n))
    dropped = {"cross_split": 0, "heldout_shingle": 0}
    keep = []
    for it in items:
        r = SPLIT_RANK[it.split]
        if rank[it.key] > r or rank[it.skey] > r:
            dropped["cross_split"] += 1
            continue
        if held_sh and not held_sh.isdisjoint(_shingles(it.skey, shingle_n)):
            dropped["heldout_shingle"] += 1
            continue
        keep.append(it)
    return keep, dropped


def _shingles(skey: str, n: int) -> list[int]:
    toks = skey.split()
    if len(toks) < n:
        return [_h(skey)] if len(toks) >= 4 else []
    return [_h(" ".join(toks[i:i + n])) for i in range(len(toks) - n + 1)]


def _h(s: str) -> int:
    return int.from_bytes(hashlib.blake2b(s.encode(), digest_size=8).digest(), "little")


# ---------------------------------------------------------------- main

def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", default="data", help="output directory (default: data)")
    ap.add_argument("--seed", type=int, default=13)
    ap.add_argument("--max-per-source", type=int, default=15000, help="record budget per in-domain source")
    ap.add_argument("--total", type=int, default=215000, help="max in-domain records (train+val+test)")
    ap.add_argument("--mix", default="0.45,0.20,0.35", help="choice,score,noul shares")
    ap.add_argument("--val-size", type=int, default=5000)
    ap.add_argument("--test-size", type=int, default=5000)
    ap.add_argument("--heldout-per-source", type=int, default=1000)
    ap.add_argument("--sources", default="", help="comma-separated subset (debug)")
    ap.add_argument("--list", action="store_true", help="list sources and exit")
    args = ap.parse_args(argv)

    if args.list:
        for n in IN_DOMAIN:
            print(f"{n:24s} {PRIMARY[n]}")
        for n in HELDOUT:
            print(f"{n:24s} {PRIMARY[n]}  (heldout)")
        return 0

    mix = dict(zip(QTYPES, (float(x) for x in args.mix.split(","))))
    only = [s for s in args.sources.split(",") if s]
    t0 = time.time()
    items: list[Item] = []
    status: dict[str, str] = {}
    for name, fn in IN_DOMAIN.items():
        if only and name not in only:
            continue
        try:
            got = run_source(name, fn, args.seed, args.max_per_source)
            items.extend(got)
            status[name] = f"ok ({len(got)} generated)"
            log(f"[{time.time() - t0:6.0f}s] {name}: {len(got)}")
        except Exception as e:  # skip sources that fail to download / parse
            status[name] = f"FAILED: {type(e).__name__}: {str(e)[:200]}"
            log(f"[{time.time() - t0:6.0f}s] {name}: FAILED\n{traceback.format_exc()}")
    held: list[Item] = []
    for name, fn in HELDOUT.items():
        if only and name not in only:
            continue
        try:
            got = dedupe(run_source(name, fn, args.seed, args.heldout_per_source))
            for it in got:
                it.split = "heldout"
            held.extend(got)
            status[name] = f"ok ({len(got)} generated, heldout)"
            log(f"[{time.time() - t0:6.0f}s] {name} (heldout): {len(got)}")
        except Exception as e:
            status[name] = f"FAILED: {type(e).__name__}: {str(e)[:200]}"
            log(f"[{time.time() - t0:6.0f}s] {name}: FAILED\n{traceback.format_exc()}")

    n_gen = len(items)
    items = dedupe(items)
    n_dedup = len(items)
    items = balance_types(items, mix, args.total, args.seed)
    n_bal = len(items)
    assign_splits(items, args.val_size, args.test_size, args.seed)
    items, dropped = remove_leaks(items, held)
    log(f"generated {n_gen}, after dedupe {n_dedup}, after type balance {n_bal}, after leak removal {len(items)} "
        f"(dropped {dropped})")

    os.makedirs(args.out, exist_ok=True)
    splits = {s: [it for it in items if it.split == s] for s in ("train", "val", "test")}
    splits["heldout"] = held
    for s, its in splits.items():
        recs = [it.rec for it in its]
        random.Random(f"{args.seed}|shuffle|{s}").shuffle(recs)
        write_jsonl(os.path.join(args.out, f"{s}.jsonl"), recs)
        log(f"wrote {s}: {len(recs)}")

    from jevlike.data.stats import write_stats
    info = {"seed": args.seed, "argv": " ".join(argv if argv is not None else sys.argv[1:]),
            "generated": n_gen, "deduped": n_dedup, "balanced": n_bal, "dropped": dropped,
            "status": status, "primary": PRIMARY, "mix": mix}
    write_stats(os.path.join(args.out, "STATS.md"), {s: [it.rec for it in its] for s, its in splits.items()}, info)
    log(f"done in {time.time() - t0:.0f}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())