"""
Data preparation for clankerDiffusion.

One streaming pass over FineWeb-edu (+ a Wikipedia slice for world knowledge,
+ synthetic tool-use and RAG/retrieval examples that teach the special tags):
  1. first N docs -> train the byte-level BPE tokenizer (from scratch)
  2. remainder    -> tokenize and pack into a flat uint16 .bin until token budget

Outputs (under --out-dir, default ./data):
  tokenizer.json  meta.json
  train.bin       (flat uint16 tokens)
  meta.json       ({n_tokens, seq_len, vocab_size})

Designed so EVERY training environment (local / Modal / Kaggle / TPU) can
regenerate the corpus itself with fast HF egress -- no 2 GB file transfer needed.
Use --scale to grow the corpus (e.g. --scale 4 for a multi-billion-token run).
"""
import os, json, random, argparse
import numpy as np
from datasets import load_dataset

import tokenizer as tokmod
from tokenizer import YKTokenizer, SPECIAL
from build_rag import FACTS, SYSTEM as RAG_SYSTEM

random.seed(1234)
np.random.seed(1234)

OUT = os.path.dirname(os.path.abspath(__file__))
DATADIR = os.path.join(OUT, "data")
os.makedirs(DATADIR, exist_ok=True)

SEQ_LEN = 1024
TOK_TRAIN_DOCS = 80_000
TOK_BUDGET_FINEWEB = 900_000_000     # scaled by --scale
TOK_BUDGET_WIKI = 300_000_000       # scaled by --scale
TOK_BUDGET_TOOL = 250_000_000       # scaled by --scale
TOK_BUDGET_RAG = 150_000_000        # scaled by --scale


# --------------------------------------------------------------------------
# 1) Tokenizer training
# --------------------------------------------------------------------------
def train_tokenizer(out_dir=DATADIR):
    print("[prep] streaming FineWeb-edu to collect tokenizer training docs ...")
    ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                      streaming=True, split="train")
    texts = []
    for i, ex in enumerate(ds):
        texts.append(ex["text"])
        if i + 1 >= TOK_TRAIN_DOCS:
            break
    print(f"[prep] collected {len(texts)} docs for tokenizer")
    tok = YKTokenizer().train(
        iter(texts), vocab_size=32768,
        save_path=os.path.join(out_dir, "tokenizer.json"))
    print(f"[prep] tokenizer trained: vocab={tok.vocab_size}")
    return tok


# --------------------------------------------------------------------------
# 2) Synthetic tool-use conversations (tag format == agent.py)
# --------------------------------------------------------------------------
CALC_TEMPLATES = [
    "What is {a} {op} {b}?", "Compute {a} {op} {b} for me.",
    "Calculate the result of {a} {op} {b}.",
    "If I start at {a} and apply {op} {b}, what do I get?",
]
OPS = {"+": "plus", "-": "minus", "*": "times", "/": "divided by"}
PY_SNIPPETS = [
    "print(sum(range(1, {n}+1)))",
    "import math\nprint(round(math.sqrt({n}), 4))",
    "print(sorted([{a}, {b}, {c}]))",
    "print({n} ** 2 + {n})",
    "s='clanker'; print(s[::-1])",
]
FILE_Q = [
    "Read the file {path} and tell me what is on the first line.",
    "What is inside {path}?", "List the files in {dir}.",
]
SYSTEM = ("You are clanker, a helpful assistant that can THINK, USE TOOLS, and "
           "USE MEMORY. You may reason in <think>...</think> at any point, "
           "interleaved with actions. Wrap tool calls in "
           "<tool name=\"...\">arguments</tool>. Available tools: calc(expr), "
           "python(code), read_file(path), list_dir(path), retrieve(query). After "
           "a tool result appears in <result>...</result>, continue and give the "
           "final answer. If <context>...</context> is provided, use it. You keep "
           "facts in a secondary memory: <mem_write>KEY<mem_kv>VALUE</mem_kv> to "
           "store, <mem_read>KEY</mem_read> to recall (result returns inside "
           "<mem_kv>...</mem_kv>), and <mem_evict>KEY</mem_evict> to forget.")


def gen_synthetic(n=60000):
    out = []
    for _ in range(n):
        kind = random.random()
        if kind < 0.45:
            a = random.randint(2, 999); b = random.randint(2, 999)
            op = random.choice(["+", "-", "*", "/"])
            b = max(2, b if op != "/" else random.randint(2, 50))
            if op == "/":
                a = a * b
            ans = eval(f"{a}{op}{b}")
            q = random.choice(CALC_TEMPLATES).format(a=a, b=b, op=OPS[op])
            tool = f'<tool name="calc">{a} {op} {b}</tool>'
            result = str(ans)
            think = f"<think>The user wants {a} {OPS[op]} {b}. I'll use the calculator.</think>"
        elif kind < 0.8:
            n_ = random.randint(3, 200); a = random.randint(1, 50); b = random.randint(1, 50); c = random.randint(1, 50)
            code = random.choice(PY_SNIPPETS).format(n=n_, a=a, b=b, c=c)
            q = f"Run this tiny Python snippet and report the output:\n{code}"
            tool = f'<tool name="python">{code}</tool>'
            try:
                import io, contextlib
                buf = io.StringIO()
                with contextlib.redirect_stdout(buf):
                    exec(code, {"__builtins__": __builtins__}, {})
                result = buf.getvalue().strip()
            except Exception as e:
                result = f"error: {e}"
            think = "<think>I can execute this with the python tool.</think>"
        else:
            path = random.choice(["notes.txt", "data/log.csv", "README.md", "config.json"])
            q = random.choice(FILE_Q).format(path=path, dir=random.choice(["src", "data", "."]))
            if "List" in q or "list" in q:
                tool = f'<tool name="list_dir">{path}</tool>'
                result = f"{path}/\n  file_a.txt\n  file_b.csv"
            else:
                tool = f'<tool name="read_file">{path}</tool>'
                result = f"line 1: hello from {path}"
            think = "<think>I should read the file with the read_file tool.</think>"

        conv = (f"<bos><system>{SYSTEM}</system>"
                f"<user>{q}</user>"
                f"<assistant>{think}{tool}<result>{result}</result>"
                f"Based on the tool result, the answer is {result}.</assistant><eos>")
        out.append(conv)
    return out


# --------------------------------------------------------------------------
# 2b) Interleaved-thinking + learned-memory examples.
#     Teaches the model to reason step-by-step *between* tool calls and to
#     persist/recall facts via <mem_write>/<mem_read>/<mem_evict> against the
#     side store (see memstore.py). Thinking is INTERLEAVED: think, act,
#     think, answer -- not one big block at the start.
# --------------------------------------------------------------------------
MEM_FACTS = [
    ("project:clanker", "clanker is a hybrid AR/diffusion LM with learned memory."),
    ("user:name", "The user's name is Ada."),
    ("user:likes", "The user likes concise answers and tool use."),
    ("fact:pi", "pi is approximately 3.14159."),
    ("fact:capitals", "The capital of France is Paris; of Japan is Tokyo."),
    ("pref:format", "Prefer <think> reasoning before tool calls."),
]
MEM_KEYS = [k for k, _ in MEM_FACTS]


def gen_memory(n=30000):
    out = []
    for _ in range(n):
        key, val = random.choice(MEM_FACTS)
        mode = random.random()
        if mode < 0.4:
            # write then read back (persistence demo)
            conv = (f"<bos><system>{SYSTEM}</system>"
                    f"<user>Remember that {val}</user>"
                    f"<assistant><think>I should store this in secondary memory "
                    f"so I can recall it later.</think>"
                    f"<mem_write>{key}<mem_kv>{val}</mem_kv>"
                    f"<think>Stored. Now I can read it back to confirm.</think>"
                    f"<mem_read>{key}</mem_read><mem_kv>{val}</mem_kv>"
                    f"Got it -- I'll remember {val}</assistant><eos>")
        elif mode < 0.75:
            # read an existing fact, interleaved with reasoning
            conv = (f"<bos><system>{SYSTEM}</system>"
                    f"<user>What do you know about {key}?</user>"
                    f"<assistant><think>Let me pull this from secondary memory.</think>"
                    f"<mem_read>{key}</mem_read><mem_kv>{val}</mem_kv>"
                    f"<think>That matches what I stored.</think> "
                    f"Based on memory: {val}</assistant><eos>")
        else:
            # evict
            conv = (f"<bos><system>{SYSTEM}</system>"
                    f"<user>Forget {key}.</user>"
                    f"<assistant><think>I'll remove it from secondary memory.</think>"
                    f"<mem_evict>{key}</mem_evict>Done, I forgot {key}.</assistant><eos>")
        out.append(conv)
    return out


# Multi-step reasoning with INTERLEAVED think/act/think/answer.
# Each entry: (question, calc_expr, final_answer)
REASON_QA = [
    ("A train travels 60 km/h for 2 hours, then 90 km/h for 1 hour. Total distance?",
     "60*2 + 90*1", "210 km"),
    ("If I buy 3 items at $4.50 each and a $2 tax, total cost?",
     "3*4.50 + 2", "$15.50"),
    ("A rectangle is 8 by 5. Area and perimeter?",
     "8*5", "area 40, perimeter 26"),
    ("Compound 5% on $1000 for 2 years?",
     "1000*1.05**2", "$1102.50"),
    ("Mix 2L at 10C with 3L at 40C, final temp?",
     "(2*10+3*40)/5", "28C"),
]
def gen_interleaved(n=30000):
    out = []
    for _ in range(n):
        q, expr, ans = random.choice(REASON_QA)
        try:
            res = str(eval(expr))
        except Exception:
            res = "?"
        conv = (f"<bos><system>{SYSTEM}</system>"
                f"<user>{q}</user>"
                f"<assistant><think>Break it into parts.</think>"
                f"<tool name=\"calc\">{expr}</tool>"
                f"<result>{res}</result>"
                f"<think>That gives the first part; combine with the rest.</think> "
                f"The answer is {ans}.</assistant><eos>")
        out.append(conv)
    return out


# --------------------------------------------------------------------------
# 3) RAG / retrieval examples (teach <tool name="retrieve"> and <context>)
# --------------------------------------------------------------------------
def gen_rag(n=40000):
    out = []
    for _ in range(n):
        topic, doc, q, a = random.choice(FACTS)
        mode = random.random()
        if mode < 0.5:
            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
                    f"<assistant><tool name=\"retrieve\">{q}</tool>"
                    f"<result>{doc}</result>{a}</assistant><eos>")
        elif mode < 0.85:
            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
                    f"<assistant><think>Let me check the provided context.</think>"
                    f"<context>{doc}</context>{a}</assistant><eos>")
        else:
            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
                    f"<assistant><think>{doc}</think>{a}</assistant><eos>")
        out.append(conv)
    return out


# --------------------------------------------------------------------------
# 4) Glaive function-calling (best-effort)
# --------------------------------------------------------------------------
def gen_glaive(max_examples=20000):
    out = []
    try:
        ds = load_dataset("glaiveai/glaive-function-calling-v2",
                          streaming=True, split="train")
    except Exception as e:
        print(f"[prep] Glaive unavailable ({e}); skipping.")
        return out
    for i, ex in enumerate(ds):
        if i >= max_examples:
            break
        try:
            conv = ex["conversations"]
            parts = ["<bos>"]
            for m in conv:
                role = m.get("role") or m.get("from")
                val = m.get("value") or m.get("content") or ""
                if role in ("system", "system_prompt"):
                    parts.append(f"<system>{val}</system>")
                elif role in ("human", "user"):
                    parts.append(f"<user>{val}</user>")
                elif role in ("gpt", "assistant", "function"):
                    val = val.replace("{\"name\":", "<tool name=\"").replace("\"function_call\"", "")
                    parts.append(f"<assistant>{val}</assistant>")
                elif role == "tool":
                    parts.append(f"<result>{val}</result>")
            parts.append("<eos>")
            out.append("".join(parts))
        except Exception:
            continue
    print(f"[prep] Glaive converted: {len(out)} examples")
    return out


# --------------------------------------------------------------------------
# 5) Packing
# --------------------------------------------------------------------------
def pack(tok, texts, bin_path, budget, seq_len):
    n = 0
    buf = []
    with open(bin_path, "ab") as f:
        for text in texts:
            ids = tok.encode(text)
            if not ids:
                continue
            buf.extend(ids)
            while len(buf) >= seq_len:
                chunk = np.array(buf[:seq_len], dtype=np.uint16)
                f.write(chunk.tobytes())
                buf = buf[seq_len:]
                n += seq_len
                if n >= budget:
                    return n
    if buf:
        chunk = np.array(buf[:seq_len], dtype=np.uint16)
        if len(chunk) == seq_len:
            with open(bin_path, "ab") as f:
                f.write(chunk.tobytes())
            n += seq_len
    return n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", default=DATADIR)
    ap.add_argument("--scale", type=float, default=1.0,
                    help="multiply token budgets (e.g. 4 -> ~4x more data)")
    ap.add_argument("--no-wiki", action="store_true")
    ap.add_argument("--no-rag", action="store_true")
    ap.add_argument("--no-glaive", action="store_true")
    ap.add_argument("--no-mem", action="store_true")
    args = ap.parse_args()

    out_dir = args.out_dir
    os.makedirs(out_dir, exist_ok=True)
    scale = args.scale
    bw_fw = int(TOK_BUDGET_FINEWEB * scale)
    bw_wiki = int(TOK_BUDGET_WIKI * scale)
    bw_tool = int(TOK_BUDGET_TOOL * scale)
    bw_rag = int(TOK_BUDGET_RAG * scale)

    tok_path = os.path.join(out_dir, "tokenizer.json")
    if os.path.exists(tok_path):
        print("[prep] loading existing tokenizer")
        tok = YKTokenizer.load(tok_path)
    else:
        # try to reuse the canonical tokenizer from HF (keeps all runs compatible)
        try:
            print("[prep] no local tokenizer; downloading canonical one from HF ...")
            from huggingface_hub import hf_hub_download
            tok_path = hf_hub_download(
                repo_id="coderofpears/clankerDiffusion-base",
                filename="data/tokenizer.json",
                repo_type="model",
                local_dir=out_dir,
                token=os.environ.get("HF_TOKEN"))
            tok = YKTokenizer.load(tok_path)
        except Exception as e:
            print(f"[prep] HF tokenizer download failed ({e}); training a new one.")
            tok = train_tokenizer(out_dir)

    bin_path = os.path.join(out_dir, "train.bin")
    if os.path.exists(bin_path):
        os.remove(bin_path)

    # --- fineweb-edu ---
    print(f"[prep] FineWeb-edu (budget {bw_fw:,}) ...")
    ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                      streaming=True, split="train")
    gen = iter(ds)
    for _ in range(TOK_TRAIN_DOCS):
        next(gen)
    def fineweb_iter():
        for ex in gen:
            yield ex["text"]
    n = pack(tok, fineweb_iter(), bin_path, bw_fw, SEQ_LEN)
    print(f"[prep] fineweb packed: {n:,} tokens")

    # --- wikipedia (world knowledge) ---
    if not args.no_wiki:
        print(f"[prep] Wikipedia (budget {bw_wiki:,}) ...")
        try:
            wds = load_dataset("wikipedia", "20220301.en",
                               streaming=True, split="train")
            def wiki_iter():
                for ex in wds:
                    yield ex["text"]
            nw = pack(tok, wiki_iter(), bin_path, bw_wiki, SEQ_LEN)
            print(f"[prep] wikipedia packed: {nw:,} tokens")
            n += nw
        except Exception as e:
            print(f"[prep] wikipedia skipped: {e}")

    # --- synthetic tool data (incl. interleaved reasoning) ---
    synth = gen_synthetic(int(60_000 * scale) + 60000)
    inter = gen_interleaved(int(30_000 * scale) + 30000)
    n2 = pack(tok, synth + inter, bin_path, bw_tool, SEQ_LEN)
    print(f"[prep] synthetic tool packed: {n2:,} tokens")

    # --- learned-memory data ---
    if not args.no_mem:
        mem = gen_memory(int(30_000 * scale) + 30000)
        nm = pack(tok, mem, bin_path, int(bw_tool * 0.6), SEQ_LEN)
        print(f"[prep] memory packed: {nm:,} tokens")
        n2 += nm

    # --- RAG / retrieval data ---
    if not args.no_rag:
        rag = gen_rag(int(40_000 * scale) + 40000)
        nr = pack(tok, rag, bin_path, bw_rag, SEQ_LEN)
        print(f"[prep] RAG packed: {nr:,} tokens")
        n2 += nr

    # --- glaive ---
    if not args.no_glaive:
        gl = gen_glaive(20000)
        n3 = pack(tok, gl, bin_path, bw_tool, SEQ_LEN) if gl else 0
    else:
        n3 = 0

    total = n + n2 + n3
    meta = {"n_tokens": int(total), "seq_len": SEQ_LEN,
            "vocab_size": tok.vocab_size, "path": "train.bin", "scale": scale}
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f)
    print(f"[prep] DONE. total tokens={total:,}  vocab={tok.vocab_size}")
    print(f"[prep] files in {out_dir}")


if __name__ == "__main__":
    main()