"""Train the GPT-U-20M tokenizer: byte-level BPE, 16,384 vocab, on ~2 GB of text sampled in the 45/35/20 mix.

Stages (finished stages are skipped on re-run):
  1. corpus   -> data/tok_corpus/<stream>.jsonl (training files) + <stream>.holdout.jsonl (holdout files)
  2. train    -> tokenizer/
  3. validate -> logs/tokenizer_report.json: fertility per domain, round-trip on 1000 samples
"""
import json
import multiprocessing as mp
import random
import time
from concurrent.futures import ProcessPoolExecutor

from tokenizers import Regex, Tokenizer, decoders, models, pre_tokenizers, trainers
from transformers import PreTrainedTokenizerFast

from common import LOG_DIR, ROOT, TOKENIZER_DIR
from sources import STREAM_DOMAIN, STREAM_SHARE, STREAMS, file_splits, open_stream

EOT = "<|endoftext|>"
VOCAB_SIZE = 16384
CORPUS_BYTES = 2_000_000_000
HOLDOUT_DOCS = 1000
CORPUS_DIR = ROOT / "data" / "tok_corpus"
SPLITS_PATH = ROOT / "data" / "file_splits.json"
# cl100k-style split: whitespace runs stay whole pre-tokens (Python indentation), digits handled by Digits below.
SPLIT_PATTERN = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def collect(name: str, files: list[str], byte_budget: float, max_docs: float, out_name: str):
    out_path = CORPUS_DIR / out_name
    tmp = out_path.with_suffix(".tmp")
    n_docs = n_bytes = 0
    t0 = time.time()
    with open(tmp, "w", encoding="utf-8") as f:
        for row in open_stream(name, files):
            f.write(json.dumps({"text": row["text"]}, ensure_ascii=False) + "\n")
            n_docs += 1
            n_bytes += len(row["text"].encode("utf-8"))
            if n_bytes >= byte_budget or n_docs >= max_docs:
                break
            if n_docs % 50_000 == 0:
                log(f"  {out_name}: {n_docs:,} docs, {n_bytes / 1e6:,.0f} MB ({n_bytes / 1e6 / (time.time() - t0):.1f} MB/s)")
    tmp.replace(out_path)
    return out_name, n_docs, n_bytes, time.time() - t0


def build_corpus(splits: dict) -> None:
    CORPUS_DIR.mkdir(parents=True, exist_ok=True)
    jobs = []
    for name in STREAMS:
        part = splits[STREAM_DOMAIN[name]]
        jobs.append((name, part["train"], STREAM_SHARE[name] * CORPUS_BYTES, float("inf"), f"{name}.jsonl"))
        jobs.append((name, part["holdout"], float("inf"), HOLDOUT_DOCS, f"{name}.holdout.jsonl"))
    pending = [j for j in jobs if not (CORPUS_DIR / j[-1]).exists()]
    log(f"corpus: {len(pending)} of {len(jobs)} files to collect")
    t0 = time.time()
    with ProcessPoolExecutor(max_workers=4, mp_context=mp.get_context("spawn")) as ex:
        for out_name, n_docs, n_bytes, secs in ex.map(collect, *zip(*pending)) if pending else []:
            log(f"  done {out_name}: {n_docs:,} docs, {n_bytes / 1e6:,.1f} MB in {secs / 60:.1f} min")
    log(f"corpus ready in {(time.time() - t0) / 60:.1f} min")


def read_texts(out_name: str, limit: int | None = None) -> list[str]:
    texts = []
    with open(CORPUS_DIR / out_name, encoding="utf-8") as f:
        for line in f:
            texts.append(json.loads(line)["text"])
            if limit and len(texts) >= limit:
                break
    return texts


def batches(out_names: list[str], size: int = 1000):
    buf = []
    for out_name in out_names:
        with open(CORPUS_DIR / out_name, encoding="utf-8") as f:
            for line in f:
                buf.append(json.loads(line)["text"])
                if len(buf) == size:
                    yield buf
                    buf = []
    if buf:
        yield buf


def build_tokenizer() -> Tokenizer:
    tok = Tokenizer(models.BPE(byte_fallback=True))
    tok.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.Split(Regex(SPLIT_PATTERN), behavior="isolated"),
        pre_tokenizers.Digits(individual_digits=True),
        pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
    ])
    tok.decoder = decoders.ByteLevel()
    return tok


def train() -> None:
    tok = build_tokenizer()
    trainer = trainers.BpeTrainer(vocab_size=VOCAB_SIZE, special_tokens=[EOT],
                                  initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)
    t0 = time.time()
    tok.train_from_iterator(batches([f"{name}.jsonl" for name in STREAMS]), trainer=trainer)
    log(f"BPE trained in {(time.time() - t0) / 60:.1f} min, vocab {tok.get_vocab_size()}")
    TOKENIZER_DIR.mkdir(parents=True, exist_ok=True)
    hf_tok = PreTrainedTokenizerFast(tokenizer_object=tok, bos_token=EOT, eos_token=EOT, pad_token=EOT,
                                     model_max_length=1024, clean_up_tokenization_spaces=False)
    hf_tok.save_pretrained(TOKENIZER_DIR)


EDGE_CASES = [
    "def f(x):\n\tif x:\n\t\treturn x\n\treturn None\n",
    "class A:\n    def __init__(self):\n        self.x = 1\n\n\n    def run(self):  # comment\n        pass\n",
    "    four\n        eight\n\t\ttabs\n  \t mixed \t\n",
    "line1\r\nline2\r\n\r\nline3\rline4",
    "Olá, ação, coração — “aspas” ‘simples’ … não é?",
    "日本語のテキスト、中文文本，한국어 텍스트",
    "emoji 👍🏽 👨‍👩‍👧‍👦 🇧🇷 ❤️",
    "مرحبا بالعالم  שלום עולם",
    "combining: e\u0301 a\u0308 n\u0303",
    "zero-width:\u200b\u200c\u200d\ufeff end",
    "numbers 1234567890 3.14159 -0.0001 1e-10 0xFF 000123",
    "   leading and trailing spaces   ",
    "\n\n\n\t\t\t\n",
    "",
    " ",
    "control \x00\x01\x1f\x7f chars",
    "math ∑∫√∞≠≤≥ αβγ ℝ²",
    "<|endoftext|> literal special token <|endoftext|>",
    "a" * 5000,
    " " * 300 + "x",
]
RANGES = [(0x20, 0x7E), (0xA0, 0x24F), (0x370, 0x3FF), (0x400, 0x4FF), (0x590, 0x6FF), (0x900, 0x97F),
          (0x3040, 0x30FF), (0x4E00, 0x4FFF), (0xAC00, 0xAD00), (0x2000, 0x206F), (0x2200, 0x22FF),
          (0x1F300, 0x1F64F), (0x1F900, 0x1F9FF)]
WHITESPACE = [" ", "  ", "    ", "\t", "\t\t", "\n", "\r\n", "\n\n", "\u00a0", "\u3000", "\n    ", "\n\t"]


def synthetic(n: int, rng: random.Random) -> list[str]:
    out = list(EDGE_CASES)
    while len(out) < n:
        parts = []
        for _ in range(rng.randint(5, 60)):
            r = rng.random()
            if r < 0.3:
                parts.append(rng.choice(WHITESPACE))
            elif r < 0.4:
                parts.append(rng.choice(EDGE_CASES))
            else:
                lo, hi = rng.choice(RANGES)
                parts.append("".join(chr(rng.randint(lo, hi)) for _ in range(rng.randint(1, 8))))
        out.append("".join(parts))
    return out[:n]


def fertility(tok: Tokenizer, texts: list[str]) -> dict:
    n_tok = sum(len(e.ids) for e in tok.encode_batch(texts))
    n_words = sum(len(t.split()) for t in texts)
    n_bytes = sum(len(t.encode("utf-8")) for t in texts)
    return {"docs": len(texts), "tokens": n_tok, "tokens_per_word": round(n_tok / n_words, 4),
            "bytes_per_token": round(n_bytes / n_tok, 4)}


def validate() -> dict:
    tok = Tokenizer.from_file(str(TOKENIZER_DIR / "tokenizer.json"))
    hf_tok = PreTrainedTokenizerFast.from_pretrained(TOKENIZER_DIR)
    held = {name: read_texts(f"{name}.holdout.jsonl") for name in STREAMS}
    n_py = round(HOLDOUT_DOCS * 0.27)
    code_mix = held["code_py"][:n_py] + held["code_other"][:HOLDOUT_DOCS - n_py]

    report = {
        "vocab_size": tok.get_vocab_size(),
        "eot_id": tok.token_to_id(EOT),
        "fits_uint16": tok.get_vocab_size() <= 65535,
        "fertility": {
            "dclm": fertility(tok, held["dclm"]),
            "edu": fertility(tok, held["edu"]),
            "code": fertility(tok, code_mix),
            "code_python_only": fertility(tok, held["code_py"]),
        },
    }

    rng = random.Random(0)
    code_all = held["code_py"] + held["code_other"]
    rng.shuffle(code_all)
    tabbed = [t for t in code_all if "\t" in t][:150]
    plain = [t for t in code_all if "\t" not in t][:300 - len(tabbed)]
    samples = held["dclm"][:300] + held["edu"][:300] + tabbed + plain + synthetic(100, rng)
    raw_fail = [i for i, s in enumerate(samples) if tok.decode(tok.encode(s).ids, skip_special_tokens=False) != s]
    hf_fail = [i for i, s in enumerate(samples)
               if hf_tok.decode(hf_tok.encode(s, add_special_tokens=False), skip_special_tokens=False,
                                clean_up_tokenization_spaces=False) != s]
    report["round_trip"] = {
        "samples": len(samples),
        "with_tabs": sum("\t" in s for s in samples),
        "non_ascii": sum(not s.isascii() for s in samples),
        "failures_tokenizers": len(raw_fail),
        "failures_transformers": len(hf_fail),
        "first_failures": [repr(samples[i][:120]) for i in (raw_fail + hf_fail)[:5]],
    }
    example = "def f(x):\n        return x\n\t\tpass  # 2024"
    report["example_tokens"] = tok.encode(example).tokens
    return report


def main() -> None:
    splits = file_splits(SPLITS_PATH)
    log("holdout files: " + ", ".join(f"{d}={v['holdout'][0]}" for d, v in splits.items()))
    build_corpus(splits)
    if (TOKENIZER_DIR / "tokenizer.json").exists():
        log("tokenizer already trained, skipping")
    else:
        train()
    report = validate()
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    (LOG_DIR / "tokenizer_report.json").write_text(json.dumps(report, indent=2, ensure_ascii=False))
    log(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()