"""Convert text or chat JSONL into streamable uint32 tokens and uint8 training masks.
Chat JSONL: {"messages": [{"role":"user", "content":"..."}, {"role":"assistant", "content":"..."}]}
Only assistant answers contribute to SFT loss. Input messages must be trusted training data.
"""
import argparse
import json
from pathlib import Path
from array import array
from tokenizers import Tokenizer


def sources(paths, fmt):
    suffix = ".txt" if fmt == "text" else ".jsonl"
    for src in paths:
        p = Path(src)
        if p.is_dir():
            yield from sorted(p.rglob("*")) if False else (x for x in sorted(p.rglob(f"*{suffix}")) if x.is_file())
        elif p.is_file():
            yield p
        else:
            raise FileNotFoundError(src)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--input", nargs="+", required=True)
    p.add_argument("--tokenizer", default="tokenizer.json")
    p.add_argument("--output-dir", required=True)
    p.add_argument("--format", choices=["text", "chat"], default="text")
    args = p.parse_args()
    tok = Tokenizer.from_file(args.tokenizer)
    eos = tok.token_to_id("<|eos|>")
    end = tok.token_to_id("<|end|>")
    if eos is None or end is None:
        raise ValueError("Tokenizer special tokens are missing")
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    count = 0
    with (out / "train.bin").open("wb") as ids_file, (out / "mask.bin").open("wb") as mask_file:
        def emit(ids, mask):
            nonlocal count
            if not ids:
                return
            array("I", ids).tofile(ids_file)
            mask_file.write(bytes([mask]) * len(ids))
            count += len(ids)

        for path in sources(args.input, args.format):
            with path.open("r", encoding="utf-8") as f:
                if args.format == "text":
                    # Read in chunks while preserving valid UTF-8 decoding.
                    while True:
                        chunk = f.read(256_000)
                        if not chunk:
                            break
                        emit(tok.encode(chunk, add_special_tokens=False).ids, 1)
                    emit([eos], 1)
                else:
                    for lineno, line in enumerate(f, 1):
                        if not line.strip():
                            continue
                        try:
                            record = json.loads(line)
                            messages = record["messages"]
                            for m in messages:
                                role = m["role"]
                                if role not in ("system", "user", "assistant"):
                                    raise ValueError(f"Unsupported role {role}")
                                rid = tok.token_to_id(f"<|{role}|>")
                                emit([rid], 0)
                                emit(tok.encode(str(m["content"]), add_special_tokens=False).ids,
                                     1 if role == "assistant" else 0)
                                emit([end], 1 if role == "assistant" else 0)
                            emit([eos], 0)
                        except (KeyError, ValueError, TypeError) as exc:
                            raise ValueError(f"{path}:{lineno}: {exc}") from exc
    if count < 3:
        raise ValueError("Too few tokens. Supply more training data.")
    (out / "metadata.json").write_text(json.dumps({"tokens": count, "format": args.format}, indent=2))
    print(f"Wrote {count:,} tokens to {out} (train.bin and mask.bin)")


if __name__ == "__main__":
    main()