#!/usr/bin/env python3
"""
Generate a tiny SYNTHETIC tool-selection dataset in DiffusionGemma format, so the
trainer/eval in this repo can be smoke-tested end-to-end WITHOUT any private data.

The real adapter was trained on private agent traces (not included). This produces
fully synthetic prompt/response pairs with the same structure: a system prompt, a
candidate tool list + a task in the user turn, the thinking-channel generation
prefill, and a dash-prefixed tool-name response ending in <turn|>.

  python3 make_example_data.py --out ./data
"""
import argparse, hashlib, json, random
from pathlib import Path

TOOLS = ["Bash", "Read", "Edit", "Write", "Grep", "Glob", "WebFetch", "WebSearch",
         "Agent", "TodoWrite", "NotebookEdit", "Task"]
SYSTEM = ("You are a tool selector. Given a task and a list of available tools, "
          "select ONLY the tools needed. Output one tool per line with a dash prefix.")
GEN_PREFILL = "<|turn>model\n<|channel>thought\n<channel|>"

# (task template, the tools it should select) — deterministic synthetic mapping
TASKS = [
    ("Read the config file at {path} and print its contents", ["Read"]),
    ("Find every TODO comment under {path} and list them", ["Grep", "Read"]),
    ("Fix the failing test in {path} — locate the bug and patch it", ["Read", "Edit", "Bash"]),
    ("Create a new module {path} with a hello function", ["Write"]),
    ("Search the web for the latest {topic} release notes", ["WebSearch", "WebFetch"]),
    ("Run the test suite and report failures", ["Bash"]),
    ("Rename the symbol {topic} across all files under {path}", ["Grep", "Edit"]),
    ("Summarize the open issues, then draft a plan", ["WebFetch", "TodoWrite"]),
    ("List all python files and count lines of code", ["Glob", "Bash"]),
    ("Delegate a deep research task about {topic}", ["Agent"]),
]
PATHS = ["src/parser.py", "lib/config.ts", "tests/test_api.py", "core/", "app/main.rs"]
TOPICS = ["MLX", "DiffusionGemma", "Rust async", "Postgres indexing", "WebGPU"]


def render(rng):
    template, tools = rng.choice(TASKS)
    task = template.format(path=rng.choice(PATHS), topic=rng.choice(TOPICS))
    # shuffle a candidate list that always includes the correct tools + distractors
    cands = list(set(tools) | set(rng.sample(TOOLS, k=rng.randint(6, 10))))
    rng.shuffle(cands)
    prompt = (f"<|turn>system\n{SYSTEM} <turn|>\n"
              f"<|turn>user\nAvailable tools: {', '.join(cands)}\n\n"
              f"Task: {task}\n\nSelect the tools needed:<turn|>\n{GEN_PREFILL}")
    response = "".join(f"- {t}\n" for t in tools).rstrip("\n") + "<turn|>"
    return {"prompt": prompt, "response": response}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="./data")
    ap.add_argument("--seed", type=int, default=7)
    args = ap.parse_args()
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    rng = random.Random(args.seed)
    for split, n in (("train", 120), ("valid", 24), ("test", 24)):
        rows = [render(rng) for _ in range(n)]
        f = out / f"{split}.jsonl"
        with open(f, "w") as fh:
            for r in rows:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"wrote {f} ({n} synthetic examples)")
    print("\nNOTE: synthetic toy data for smoke-testing the pipeline only — not the "
          "real training corpus. Expect the model to overfit this tiny set quickly.")


if __name__ == "__main__":
    main()