#!/usr/bin/env python3
"""
Build a leakage-free DiffusionGemma tool-selector dataset.

The original mlx_selector splits are 96% contaminated: 3,935 rows contain only
~823 distinct (user, assistant) pairs, and 299/311 unique test pairs appear
verbatim in train — any eval on that split measures memorization (a trivial
train-lookup baseline scores Jaccard 0.959).

This script:
  1. merges all mlx_selector splits and dedups exact (user, assistant) pairs
  2. groups pairs by task identity (task-text prefix after "Task:") so
     near-duplicate prompts from the same workflow cannot straddle splits
  3. greedy-packs groups into 70/15/15 train/valid/test with zero overlap
  4. verifies: no (user, assistant) pair overlap AND no user-prompt overlap
  5. renders to DiffusionGemma chat format (same as build_diffusiongemma_data:
     trailing space in system turn to byte-match mlx-vlm serving, model-turn
     thought-channel prefill, response + <turn|>)
  6. writes manifest.json: source sha256s, counts, group stats, split sha256s

Output: processed/diffusiongemma_clean/{train,valid,test}.jsonl + manifest.json
"""

import argparse
import hashlib
import json
from collections import defaultdict
from pathlib import Path

import random

GEN_PREFILL = "<|turn>model\n<|channel>thought\n<channel|>"


def sha256_file(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def group_key(user: str) -> str:
    """Task-identity key: first 160 chars of the task text (after the
    candidate tool list), so same-workflow near-duplicates group together."""
    marker = "\n\nTask: "
    i = user.find(marker)
    task = user[i + len(marker):] if i >= 0 else user
    return hashlib.sha256(task[:160].encode()).hexdigest()


def render(sys_msg, usr_msg, ast_msg):
    parts = []
    if sys_msg:
        # trailing space byte-matches mlx-vlm's apply_chat_template (serve path)
        parts.append(f"<|turn>system\n{sys_msg.strip()} <turn|>\n")
    parts.append(f"<|turn>user\n{usr_msg.strip()}<turn|>\n")
    parts.append(GEN_PREFILL)
    return {"prompt": "".join(parts), "response": f"{ast_msg.strip()}<turn|>"}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default="./mlx_selector")
    ap.add_argument("--dst", default="./diffusiongemma_clean")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--ratios", default="0.70,0.15,0.15")
    args = ap.parse_args()
    src, dst = Path(args.src), Path(args.dst)
    dst.mkdir(parents=True, exist_ok=True)

    # 1. merge + dedup exact pairs
    pairs = {}  # (user, assistant) -> (sys, user, assistant)
    src_hashes, total_rows = {}, 0
    for split in ("train", "valid", "test"):
        f = src / f"{split}.jsonl"
        src_hashes[f.name] = sha256_file(f)
        with open(f) as fh:
            for line in fh:
                obj = json.loads(line)
                m = {x["role"]: x["content"] for x in obj["messages"]}
                total_rows += 1
                pairs[(m["user"], m["assistant"])] = (m.get("system", ""), m["user"], m["assistant"])

    # 2. group by task identity
    groups = defaultdict(list)
    for (user, _ast), triple in pairs.items():
        groups[group_key(user)].append(triple)
    group_items = sorted(groups.items())  # deterministic order
    rng = random.Random(args.seed)
    rng.shuffle(group_items)

    # 3. greedy pack into splits by pair count
    ratios = [float(x) for x in args.ratios.split(",")]
    targets = [r * len(pairs) for r in ratios]
    buckets = [[], [], []]
    counts = [0, 0, 0]
    for gk, triples in group_items:
        # assign to the split furthest below its target (relative)
        deficits = [(counts[i] / targets[i] if targets[i] else 1.0, i) for i in range(3)]
        i = min(deficits)[1]
        buckets[i].extend(triples)
        counts[i] += len(triples)

    names = ["train", "valid", "test"]
    # 4. verify zero overlap
    pair_sets = [set((u, a) for _, u, a in b) for b in buckets]
    user_sets = [set(u for _, u, _ in b) for b in buckets]
    for i in range(3):
        for j in range(i + 1, 3):
            assert not (pair_sets[i] & pair_sets[j]), f"pair overlap {names[i]}/{names[j]}"
            assert not (user_sets[i] & user_sets[j]), f"prompt overlap {names[i]}/{names[j]}"

    # 5. render + write
    split_stats = {}
    for name, bucket in zip(names, buckets):
        out = dst / f"{name}.jsonl"
        rng2 = random.Random(args.seed + 7)
        bucket = list(bucket)
        rng2.shuffle(bucket)
        with open(out, "w") as fh:
            for sys_msg, usr, ast in bucket:
                fh.write(json.dumps(render(sys_msg, usr, ast), ensure_ascii=False) + "\n")
        split_stats[name] = {"pairs": len(bucket), "sha256": sha256_file(out)}

    # 6. manifest
    manifest = {
        "source": {"dir": str(src), "files": src_hashes, "total_rows": total_rows},
        "distinct_pairs": len(pairs),
        "groups": len(groups),
        "seed": args.seed,
        "ratios": ratios,
        "splits": split_stats,
        "note": "deduped exact (user,assistant) pairs; group-aware split by task-text prefix (160 chars); zero pair AND zero prompt overlap verified; ~44% of prompts carry an upstream ~2990-char truncation from the original mlx_selector build",
    }
    (dst / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()