import time
import datetime
import os
import gc
import json
import mmap as mmap_mod
import pickle
import unicodedata
import importlib.resources
import requests
import numpy as np
import regex
import re
from .tokenizer import (
    KinTokenizer,
    _HAS_RUST_MERGER,
    _save_token_chunks,
    _load_token_chunks,
    _token_chunks_path,
    _TOKEN_CHUNKS_LEGACY_EXT,
)
from .params import (
    METASPACE,
    DEFAULT_VOCAB_SIZE,
    MIN_VOCAB_SIZE,
    DEFAULT_MIN_FREQUENCY,
    DEFAULT_STREAM_LINES,
    RUST_MERGE_BATCH_SIZE,
    CHECKPOINT_EVERY_N_MERGES,
    LOAN_WORD_ALLOWLIST,
    MMAP_THRESHOLD_BYTES,
    DEFAULT_EVAL_RATIO,
    DEFAULT_EVAL_SEED,
)

try:
    from kin_merge import (
        rust_preprocess_text as _rust_preprocess,
        rust_stream_pretokenize_file as _rust_stream_pretokenize_file,
        rust_filter_words as _rust_filter_words,
        rust_create_tokens_parallel as _rust_create_tokens_parallel,
        rust_bpe_train_loop as _rust_bpe_train_loop,
        rust_build_bpe_trainer as _rust_build_bpe_trainer,
        rust_build_bpe_trainer_from_file as _rust_build_bpe_trainer_from_file,
        rust_encode_text_raw as _rust_encode_text_raw,
    )
except Exception:
    _rust_preprocess = None
    _rust_stream_pretokenize_file = None
    _rust_filter_words = None
    _rust_create_tokens_parallel = None
    _rust_bpe_train_loop = None
    _rust_build_bpe_trainer = None
    _rust_build_bpe_trainer_from_file = None
    _rust_encode_text_raw = None

try:
    from kin_merge import rust_encode_text_raw_batch as _rust_encode_text_raw_batch
except Exception:
    _rust_encode_text_raw_batch = None


def _format_duration(total_seconds):
    """Format a duration in seconds to a human-readable string."""
    months, remaining = int(total_seconds // (3600 * 24 * 30)), total_seconds % (3600 * 24 * 30)
    days, remaining = int(remaining // (3600 * 24)), remaining % (3600 * 24)
    hours, remaining = int(remaining // 3600), remaining % 3600
    minutes, seconds = int(remaining // 60), int(remaining % 60)
    parts = []
    if months > 0:
        parts.append(f"{months} months")
    if days > 0:
        parts.append(f"{days} days")
    if hours > 0:
        parts.append(f"{hours} hours")
    if minutes > 0:
        parts.append(f"{minutes} minutes")
    if seconds > 0 or not parts:
        parts.append(f"{seconds} seconds")
    return " ".join(parts)


def _configure_rayon_threads(nbr_processes):
    """Configure Rayon worker count before the first parallel Rust call."""
    if nbr_processes is None:
        return
    if nbr_processes < 1:
        raise ValueError("nbr_processes must be at least 1")
    os.environ["RAYON_NUM_THREADS"] = str(int(nbr_processes))



def train_kin_tokenizer(text, vocab_size=DEFAULT_VOCAB_SIZE, save=False, tokenizer_path=None, retrain=False, allow_legacy_pickle=False, nbr_processes=None, lowercase=None):
    """
    Function for training the tokenizer
    params:
        text: the string text that will be used for training the tokenizer
        vocab_size: the final size of the vocabulary for the tokenizer
        save: boolean to indicate if tokenizer has to be saved after training for future use
        tokenizer_path: the path to which the tokenizer will be saved if save is True
        retrain: continue training from checkpoint at tokenizer_path
        allow_legacy_pickle: allow loading/saving legacy pickle checkpoint (unsafe)
        nbr_processes: number of worker processes for training stages
        lowercase: whether to lowercase during tokenizer training; None preserves checkpoint behavior on retrain
    Returns:
        returns tokenizer object after training
    """
    tokenizer = KinTokenizer(autoload=False)
    start_merge_iter = 1
    if retrain:
        if tokenizer_path is None:
            raise ValueError("tokenizer_path is required when retrain=True")
        tokenizer.load(tokenizer_path, allow_pickle=allow_legacy_pickle)
        start_merge_iter = _infer_start_merge_iter(tokenizer)
    if len(text) < vocab_size or not isinstance(text, str):
        raise ValueError("length of text should be greater or equal to vocab_size, vocab_size should be at least 256 and text should be a string")
    
    if save == True:
        if tokenizer_path is None:
           tokenizer_path = os.path.join("kin_tokenizer", "data")
        
        tokenizer.train(
            text,
            vocab_size,
            start_merge_iter=start_merge_iter,
            tokenizer_path=tokenizer_path,
            nbr_processes=nbr_processes,
            lowercase=lowercase,
        )
        tokenizer.save(tokenizer_path, save_legacy_pickle=allow_legacy_pickle)
    else:
        tokenizer.train(
            text,
            vocab_size,
            start_merge_iter=start_merge_iter,
            nbr_processes=nbr_processes,
            lowercase=lowercase,
        )

    return tokenizer


def train_kin_tokenizer_streamed_from_file(
    dataset_path,
    vocab_size=DEFAULT_VOCAB_SIZE,
    save=False,
    tokenizer_path=None,
    retrain=False,
    allow_legacy_pickle=False,
    stream_lines=DEFAULT_STREAM_LINES,
    min_frequency=DEFAULT_MIN_FREQUENCY,
    verbose=True,
    lowercase=None,
    nbr_processes=None,
    save_token_chunks=True,
):
    """
    Rust-first streamed training path:
    - stream dataset loading + pretokenization in Rust
    - keep filtering/token creation in Rust
    - run merge loop in Rust batches

    This path avoids building one huge Python text string for large corpora.
    """
    if not _HAS_RUST_MERGER or _rust_stream_pretokenize_file is None:
        raise RuntimeError("Rust extension is required for streamed training path")

    if dataset_path is None or not os.path.isfile(dataset_path):
        raise ValueError("dataset_path must point to an existing text file")

    tokenizer = KinTokenizer(autoload=False)
    start_merge_iter = 1

    if retrain:
        if tokenizer_path is None:
            raise ValueError("tokenizer_path is required when retrain=True")
        tokenizer.load(tokenizer_path, allow_pickle=allow_legacy_pickle)
        start_merge_iter = _infer_start_merge_iter(tokenizer)
    else:
        tokenizer.reset_for_training()

    lowercase = tokenizer._should_lowercase(lowercase)
    tokenizer.normalization["lowercase"] = lowercase
    _configure_rayon_threads(nbr_processes)

    if tokenizer_path is None:
        tokenizer_path = os.path.join("kin_tokenizer", "data")
    tokenizer_path = os.path.abspath(tokenizer_path)
    os.makedirs(tokenizer_path, exist_ok=True)

    if vocab_size < MIN_VOCAB_SIZE:
        raise ValueError(f"vocab_size must be at least {MIN_VOCAB_SIZE}")

    print(f"\n{'=' * 60}")
    print(f"  Streamed BPE Training Configuration")
    print(f"{'=' * 60}")
    print(f"  Dataset:          {dataset_path}")
    print(f"  Target vocab:     {vocab_size}")
    print(f"  Stream lines:     {stream_lines}")
    print(f"  Min frequency:    {min_frequency}")
    print(f"  Lowercase:        {lowercase}")
    print(f"  Save chunks:      {save_token_chunks}")
    print(f"  Checkpoint dir:   {tokenizer_path}")
    print(f"  Retrain:          {retrain} (start_merge_iter={start_merge_iter})")
    print(f"{'=' * 60}\n")

    trainer = None
    tokens_chunks = None
    next_token_id = None
    if start_merge_iter < 2:
        tokenizer.ensure_base_byte_vocab()
        next_token_id = max(tokenizer._KinTokenizer__vocab.keys()) + 1
        non_kinyarwanda_words_json_obj = importlib.resources.files('kin_tokenizer').joinpath('data/non_kinyarwanda_words.json')
        try:
            with importlib.resources.as_file(non_kinyarwanda_words_json_obj) as json_path:
                with open(json_path, "r", encoding="utf-8") as f:
                    non_kinyarwanda_words = json.load(f)
        except FileNotFoundError:
            non_kinyarwanda_words_path_obj = importlib.resources.files('kin_tokenizer').joinpath('data/non_kinyarwanda_words.pkl')
            with importlib.resources.as_file(non_kinyarwanda_words_path_obj) as non_kinyarwanda_words_file_path:
                with open(non_kinyarwanda_words_file_path, "rb") as f:
                    non_kinyarwanda_words = pickle.load(f)

        # FIX #18: Protect common French/English loan words used in Kinyarwanda.
        filtered_non_kin = [w for w in non_kinyarwanda_words if w.lower() not in LOAN_WORD_ALLOWLIST]
        existing_words_in_vocab = [
            value.decode("UTF-8", errors="ignore")
            for _, value in tokenizer._KinTokenizer__vocab.items()
            if isinstance(value, bytes)
        ]
        if _rust_build_bpe_trainer_from_file is not None:
            print("[Step 1/4] Streaming pretokenization + tokenization + trainer build in Rust...")
            trainer = _rust_build_bpe_trainer_from_file(
                dataset_path,
                next_token_id,
                filtered_non_kin,
                existing_words_in_vocab,
                stream_lines,
                True,
                lowercase,
                METASPACE,
            )
            del filtered_non_kin
            print(f"[Step 1/4] Trainer built. Total tokens: {trainer.total_tokens:,}")
        else:
            print("[Step 1/4] Streaming pretokenization from file in Rust...")
            chunks = _rust_stream_pretokenize_file(dataset_path, stream_lines, True, lowercase)

            if _rust_filter_words is not None:
                # Chunks have ▁ prefix (metaspace); strip it before comparing against the
                # bare-word filter list so "▁Aaron" is correctly matched against "Aaron".
                chunks = _rust_filter_words(chunks, filtered_non_kin, METASPACE)
                del filtered_non_kin

            if _rust_filter_words is not None:
                chunks = _rust_filter_words(chunks, existing_words_in_vocab)

            print("[Step 2/4] Creating byte-token chunks in Rust...")
            if _rust_create_tokens_parallel is None:
                raise RuntimeError("rust_create_tokens_parallel is unavailable")
            tokens_chunks = _rust_create_tokens_parallel(chunks)
            del chunks
            gc.collect()
            if _rust_build_bpe_trainer is not None:
                trainer = _rust_build_bpe_trainer(tokens_chunks, next_token_id)
                if not save_token_chunks:
                    tokens_chunks = None
                    gc.collect()
    else:
        last_token_chunks_file_path = _token_chunks_path(tokenizer_path)
        legacy_pkl_path = os.path.join(tokenizer_path, "current_token_chunks" + _TOKEN_CHUNKS_LEGACY_EXT)
        if not os.path.isfile(last_token_chunks_file_path):
            if os.path.isfile(legacy_pkl_path):
                raise ValueError(
                    f"Only legacy pickle chunk file found: {legacy_pkl_path}. "
                    "Pickle token chunks are no longer loaded due to security risks. "
                    "Re-run first-train (without --retrain) to generate a safe .bin checkpoint."
                )
            raise ValueError(
                f"Missing token chunk checkpoint for retrain: {last_token_chunks_file_path}. "
                "Run non-streamed retrain once or provide first-train mode."
            )
        tokens_chunks = _load_token_chunks(last_token_chunks_file_path)
        next_token_id = max(tokenizer._KinTokenizer__vocab.keys()) + 1
        if _rust_build_bpe_trainer is not None:
            trainer = _rust_build_bpe_trainer(tokens_chunks, next_token_id)
            if not save_token_chunks:
                tokens_chunks = None
                gc.collect()

    num_merges = vocab_size - len(tokenizer._KinTokenizer__vocab)
    if num_merges <= 0:
        tokenizer.vocab_size = len(tokenizer._KinTokenizer__vocab)
        if save:
            tokenizer.save(tokenizer_path, save_legacy_pickle=allow_legacy_pickle)
        return tokenizer

    if next_token_id is None:
        next_token_id = max(tokenizer._KinTokenizer__vocab.keys()) + 1
    merges_remaining = num_merges
    total_merges_done = 0
    # FIX #8: Increased batch size to reduce Python↔Rust boundary copies.
    rust_batch_size = RUST_MERGE_BATCH_SIZE
    # FIX #19: Use min_frequency directly. The previous formula
    # (total_tokens // 1_000_000) produced values like 1019 that stopped
    # training far short of the requested vocab_size.
    if trainer is not None:
        total_tokens_approx = trainer.total_tokens
    elif tokens_chunks is not None:
        total_tokens_approx = sum(len(c) for c in tokens_chunks)
    else:
        total_tokens_approx = 0
    train_min_frequency = min_frequency
    log_file_path = os.path.join(tokenizer_path, "training_logs.txt")
    last_token_chunks_file_path = _token_chunks_path(tokenizer_path)

    if trainer is None and _rust_bpe_train_loop is None:
        raise RuntimeError("rust_bpe_train_loop is unavailable")

    loop_start = time.time()
    with open(log_file_path, "a" if start_merge_iter > 1 else "w", encoding="utf-8") as f:
        rust_threads = os.environ.get("RAYON_NUM_THREADS", str(os.cpu_count() or 1))
        header = (
            f"\n{'=' * 60}\n"
            f"  [Step 3/4] BPE Merge Loop\n"
            f"{'=' * 60}\n"
            f"  Rust worker threads:  {rust_threads}\n"
            f"  BPE min_frequency:    {train_min_frequency}\n"
            f"  Total tokens:         {total_tokens_approx:,}\n"
            f"  Merges requested:     {num_merges}\n"
            f"  Current vocab size:   {len(tokenizer._KinTokenizer__vocab)}\n"
            f"  Target vocab size:    {vocab_size}\n"
            f"  Batch size:           {rust_batch_size}\n"
            f"  Save token chunks:    {save_token_chunks}\n"
            f"{'=' * 60}"
        )
        print(header)
        f.write(header + "\n")

        while merges_remaining > 0:
            batch_size = min(rust_batch_size, merges_remaining)
            if trainer is not None:
                merge_rules = trainer.step(batch_size, train_min_frequency)
            else:
                trained_chunks, merge_rules, _ = _rust_bpe_train_loop(
                    tokens_chunks,
                    next_token_id,
                    batch_size,
                    train_min_frequency,
                )
                tokens_chunks = trained_chunks

            actual_batch = len(merge_rules)
            if actual_batch == 0:
                early_msg = (
                    f"\nMerge loop stopped: no pairs with frequency >= {train_min_frequency}. "
                    f"Completed {total_merges_done}/{num_merges} merges. "
                    f"Final vocab size: {len(tokenizer._KinTokenizer__vocab)}."
                )
                print(early_msg)
                f.write(early_msg + "\n")
                break

            for merge_offset, merge_rule in enumerate(merge_rules):
                global_idx = start_merge_iter + total_merges_done + merge_offset
                top_pair, new_token = merge_rule
                tokenizer.merged_tokens[top_pair] = new_token
                tokenizer._KinTokenizer__vocab[new_token] = (
                    tokenizer._KinTokenizer__vocab[top_pair[0]] + tokenizer._KinTokenizer__vocab[top_pair[1]]
                )

                if verbose:
                    line = (
                        f"Merge({global_idx}/{num_merges + start_merge_iter - 1}): "
                        f"Pair {top_pair} -> {new_token}. Vocab size: {len(tokenizer._KinTokenizer__vocab)}"
                    )
                    print(line)
                    f.write(line + "\n")

            total_merges_done += actual_batch
            merges_remaining -= actual_batch
            next_token_id += actual_batch

            # PERF FIX: Only checkpoint every N merges instead of every batch.
            if save and (total_merges_done % CHECKPOINT_EVERY_N_MERGES == 0 or merges_remaining == 0):
                tokenizer.vocab_size = len(tokenizer._KinTokenizer__vocab)
                tokenizer.save(tokenizer_path, save_legacy_pickle=allow_legacy_pickle)
                if save_token_chunks:
                    checkpoint_chunks = trainer.snapshot_chunks() if trainer is not None else tokens_chunks
                    _save_token_chunks(checkpoint_chunks, last_token_chunks_file_path)
                    del checkpoint_chunks

            gc.collect()

        summary_line = (
            f"\n{'=' * 60}\n"
            f"  [Step 3/4] Merge Loop Complete\n"
            f"{'=' * 60}\n"
            f"  Total merges done:    {total_merges_done}/{num_merges}\n"
            f"  Final vocab size:     {len(tokenizer._KinTokenizer__vocab)}\n"
            f"  Time elapsed:         {(time.time() - loop_start) // 60:.0f} min {round((time.time() - loop_start) % 60)} seconds\n"
            f"{'=' * 60}"
        )
        print(summary_line)
        f.write(summary_line + "\n")

    print("[Step 4/4] Adding special tokens...")
    tokenizer.ensure_special_tokens()
    tokenizer.vocab_size = len(tokenizer._KinTokenizer__vocab)
    tokenizer._reset_token_lookup_cache()
    print(f"Special tokens added. Final vocabulary size: {tokenizer.vocab_size}")

    if save:
        tokenizer.save(tokenizer_path, save_legacy_pickle=allow_legacy_pickle)

    return tokenizer


def _infer_start_merge_iter(tokenizer):
    """
    Infer next merge iteration from already learned merge rules.
    Iteration numbering is merge-count based, not token-id based.
    """
    return max(1, len(tokenizer.merged_tokens) + 1)


# FIX #5: fix_case removed — vowel-capital regex no longer needed.

# Pre-compiled URL pattern — avoids re.compile() on every preprocess_text call.
_RE_URL = re.compile(r'https?://\S+|www\.\S+')

def clean_lines(text):
    """
    Remove lines that contain only numbers, special characters,
    or mixes of numbers and special characters (no letters).
    """
    return '\n'.join(
        line for line in text.splitlines()
        if not regex.fullmatch(r'\s*([^\w\s]*\d+[^\w\s]*|[^\w\s]+|\d+)\s*', line)
    )

def preprocess_text(text, is_lowercase_text=False):
    """Preprocess text for dataset creation. Mirrors the preprocessing in KinTokenizer.train().
    
    Prefers the Rust-accelerated path when the extension is available — this is
    3-5× faster than Python regex for large corpora and avoids creating 12+
    intermediate copies of the full text string.
    """
    if _rust_preprocess is not None:
        # Rust preprocess_text_inner already applies NFC internally,
        # so skip the redundant Python NFC pass (~2s + 1.8 GB copy saved).
        text = _rust_preprocess(text, bool(is_lowercase_text))
    else:
        # Python fallback: NFC first, then regex cleanup.
        text = unicodedata.normalize("NFC", text)

        text = regex.sub("\u2018", "'", text)
        text = regex.sub("\u2019", "'", text)
        text = regex.sub("\u201c", '"', text)
        text = regex.sub("\u201d", '"', text)
       
        text = clean_lines(text)
        text = regex.sub(r'(\n){3,}', '\n\n', text)

        text = regex.sub(r'[^\S\n]+\n', '\n', text)
        text = regex.sub(r'^(?!\s*$)\s+', '', text, flags=regex.MULTILINE)
        # FIX: Remove URLs BEFORE punctuation-joining to match Rust order,
        # preventing URL colons/dots from being mangled by the punct regex.
        text = _RE_URL.sub('', text)
        text = regex.sub(r'\*', ' ', text)
        text = regex.sub(r'[^\S\r\n]+', ' ', text)

        text = regex.sub(r'(\w)\s+([.,;:!?])', r'\1\2', text)

        if is_lowercase_text:
            text = text.lower()

    return text

def _read_text_efficient(path):
    """Read a large UTF-8 text file with lower peak memory via mmap for big files."""
    file_size = os.path.getsize(path)
    if file_size >= MMAP_THRESHOLD_BYTES:
        with open(path, "rb") as f:
            with mmap_mod.mmap(f.fileno(), 0, access=mmap_mod.ACCESS_READ) as mm:
                return mm.read().decode("UTF-8", errors="ignore")
    with open(path, "r", encoding="UTF-8", errors="ignore") as f:
        return f.read()


def _split_paragraphs_train_eval(text, eval_ratio, seed):
    """Split preprocessed text into train/eval at paragraph boundaries.

    Paragraphs are delimited by double newlines (``\\n\\n``).  They are
    shuffled deterministically with *seed*, then greedily assigned to
    eval until *eval_ratio* of the total character budget is reached.
    Remaining paragraphs go to train.

    After assignment, each split's paragraphs are **sorted back to their
    original document order** so that consecutive paragraphs stay coherent
    for overlapping-window sequence creation.

    Returns ``(train_text, eval_text)``.
    """
    if eval_ratio <= 0:
        return text, ""
    if eval_ratio >= 1.0:
        return "", text

    paragraphs = [p for p in text.split("\n\n") if p.strip()]
    if len(paragraphs) < 2:
        return text, ""

    rng = np.random.default_rng(seed)
    indices = np.arange(len(paragraphs))
    rng.shuffle(indices)

    total_chars = sum(len(paragraphs[i]) for i in indices)
    eval_target = int(total_chars * eval_ratio)

    eval_indices = []
    train_indices = []
    eval_chars = 0

    for idx in indices:
        if eval_chars < eval_target:
            eval_indices.append(int(idx))
            eval_chars += len(paragraphs[int(idx)])
        else:
            train_indices.append(int(idx))

    # Restore original document order within each split so that
    # overlapping sliding-window sequences see locally coherent text.
    train_indices.sort()
    eval_indices.sort()

    train_text = "\n\n".join(paragraphs[i] for i in train_indices)
    eval_text = "\n\n".join(paragraphs[i] for i in eval_indices)

    return train_text, eval_text


def _numpy_stride_sequences(tokens_array, sequence_length, step_size):
    """Create overlapping sequences using numpy stride tricks — zero-copy views.

    Instead of slicing N separate Python lists (each allocating new memory),
    this creates a single strided *view* of the contiguous token array.
    The view shares memory with the source — no data is copied until the
    final write to the mmap file.

    Memory saving: for 900M tokens, seq_len=512, step=256:
      - Old (list-of-lists): ~1.75M × 513 × 28B (Python int) ≈ 25 GB peak
      - New (stride view):   ~3.6 GB (the one np.int32 array) + tiny view header
    """
    total = len(tokens_array)
    seq_width = sequence_length + 1
    if total < seq_width or step_size == 0:
        return np.empty((0, seq_width), dtype=tokens_array.dtype)
    n_seqs = (total - seq_width) // step_size + 1
    stride_bytes = tokens_array.strides[0]
    return np.lib.stride_tricks.as_strided(
        tokens_array,
        shape=(n_seqs, seq_width),
        strides=(step_size * stride_bytes, stride_bytes),
    )


def _encode_text_to_tokens(text, tokenizer, nbr_processes, label=""):
    """Encode preprocessed+normalized text into a contiguous np.int32 token array.

    Handles Rust batch / Rust per-chunk / Python fallback paths automatically.
    Returns the token array and the encoding wall-clock time in seconds.
    """
    phase_start = time.time()
    prefix = f"  [{label}] " if label else "  "

    _use_rust_raw = (
        _rust_encode_text_raw is not None
        and getattr(tokenizer, "space_strategy", "metaspace") == "metaspace"
    )
    _use_rust_raw_batch = _use_rust_raw and _rust_encode_text_raw_batch is not None

    # ── Split text into ~50 MB chunks at word boundaries ──
    ENCODE_CHUNK_BYTES = 50 * 1024 * 1024
    text_len = len(text)
    chunks = []
    pos = 0
    while pos < text_len:
        end = min(pos + ENCODE_CHUNK_BYTES, text_len)
        if end < text_len:
            nl = text.rfind('\n', pos, end)
            if nl > pos:
                end = nl + 1
            else:
                sp = text.rfind(' ', pos, end)
                if sp <= pos:
                    sp = text.rfind('\u2581', pos, end)
                if sp > pos:
                    end = sp
        chunks.append(text[pos:end])
        pos = end

    estimated_tokens = text_len // 3 + 1024
    tokens_array = np.empty(estimated_tokens, dtype=np.int32)
    write_pos = 0

    if _use_rust_raw_batch:
        raw = _rust_encode_text_raw_batch(chunks, tokenizer.merged_tokens)
        del chunks
        arr = np.frombuffer(raw, dtype=np.int32)
        total_tokens = len(arr)
        if total_tokens > estimated_tokens:
            tokens_array = np.empty(total_tokens, dtype=np.int32)
        tokens_array[:total_tokens] = arr
        write_pos = total_tokens
        del raw, arr
        elapsed = time.time() - phase_start
        print(f"{prefix}Batch encoded: {total_tokens:,} tokens ({elapsed:.1f}s)")
    else:
        for chunk_idx, chunk in enumerate(chunks):
            if _use_rust_raw:
                raw = _rust_encode_text_raw(chunk, tokenizer.merged_tokens)
                arr = np.frombuffer(raw, dtype=np.int32)
                n = len(arr)
                if write_pos + n > len(tokens_array):
                    new_size = max(len(tokens_array) * 2, write_pos + n)
                    tokens_array = np.resize(tokens_array, new_size)
                tokens_array[write_pos:write_pos + n] = arr
                write_pos += n
                del raw, arr
            else:
                token_ids = tokenizer.encode(chunk, nbr_processes=nbr_processes, _skip_normalize=True)
                n = len(token_ids)
                if write_pos + n > len(tokens_array):
                    new_size = max(len(tokens_array) * 2, write_pos + n)
                    tokens_array = np.resize(tokens_array, new_size)
                tokens_array[write_pos:write_pos + n] = token_ids
                write_pos += n
                del token_ids

            del chunk
            if (chunk_idx + 1) % 5 == 0 or chunk_idx == len(chunks) - 1:
                elapsed = time.time() - phase_start
                pct = (chunk_idx + 1) / len(chunks) * 100
                print(f"{prefix}Chunk {chunk_idx + 1}: {write_pos:,} tokens so far ({pct:.0f}%, {elapsed:.1f}s)")

        del chunks
        gc.collect()

    tokens_array = tokens_array[:write_pos]
    encode_time = time.time() - phase_start
    return tokens_array, encode_time


def _write_sequences_npy(tokens_array, sequence_length, step_size, file_path, label=""):
    """Create stride-tricks sequences from *tokens_array* and write to *file_path*.

    Returns ``(total_sequences, write_time_seconds)``.
    """
    prefix = f"[{label}] " if label else ""
    seq_width = sequence_length + 1

    if len(tokens_array) < seq_width:
        print(f"{prefix}Skipping — only {len(tokens_array):,} tokens (need {seq_width})")
        return 0, 0.0

    phase_start = time.time()
    seq_view = _numpy_stride_sequences(tokens_array, sequence_length, step_size)
    total_sequences = seq_view.shape[0]
    if total_sequences == 0:
        print(f"{prefix}No sequences generated")
        return 0, 0.0
    print(f"{prefix}{total_sequences:,} sequences created via stride view ({time.time() - phase_start:.1f}s)")

    phase_start = time.time()
    mmap_array = np.lib.format.open_memmap(
        file_path, mode='w+', dtype=np.int32, shape=(total_sequences, seq_width)
    )
    WRITE_BATCH = 100_000
    for batch_start in range(0, total_sequences, WRITE_BATCH):
        batch_end = min(batch_start + WRITE_BATCH, total_sequences)
        mmap_array[batch_start:batch_end] = seq_view[batch_start:batch_end].copy()

    del seq_view, mmap_array
    gc.collect()

    write_time = time.time() - phase_start
    file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
    print(f"{prefix}Saved to {file_path} ({file_size_mb:.1f} MB, {write_time:.1f}s)")
    return total_sequences, write_time


def create_dataset(text_file_path: str, nbr_processes: int, sequence_length: int, destination_dir: str, step_size:int=None, is_text_file_path_url:bool=False, is_lowercase_text:bool=None, tokenizer_path:str=None, eval_ratio:float=DEFAULT_EVAL_RATIO, eval_seed:int=DEFAULT_EVAL_SEED):
    """Create train (and optionally eval) sequence files for language model training.

    This only affects downstream language-model sequences. It does NOT split the
    corpus used to train or retrain the tokenizer vocabulary itself.

    The text is split at **paragraph level** (double-newline boundaries) BEFORE
    encoding so that train and eval sequences have **zero token overlap**.
    Paragraphs are shuffled deterministically with *eval_seed* before splitting
    to ensure the eval set is a representative sample of the full corpus.

    Outputs:
        * ``train_sequences.npy`` — the main training sequences
        * ``eval_sequences.npy``  — held-out evaluation sequences (when eval_ratio > 0)

    When *eval_ratio* is 0, only ``train_sequences.npy`` is produced (equivalent
    to the previous single-file behaviour).

    params:
        text_file_path: filename path or URL containing the text
        nbr_processes: number of parallel processes for encoding
        sequence_length: tokens per sequence (+1 target token added automatically)
        destination_dir: directory where .npy files are saved
        step_size: sliding-window stride; None → sequence_length // 2
        is_text_file_path_url: treat text_file_path as a URL
        is_lowercase_text: lowercase text before creating sequences;
            None uses the loaded tokenizer setting
        tokenizer_path: directory or file path to load the tokenizer from;
            None uses KinTokenizer() auto-load
        eval_ratio: fraction of text (by paragraph character count) held out
            for evaluation (default: 0.05 = 5%)
        eval_seed: random seed for deterministic paragraph shuffle (default: 42)
    """
    training_start_time = datetime.datetime.now(datetime.timezone.utc)
    start_time = time.time()

    if tokenizer_path is not None:
        tokenizer = KinTokenizer(autoload=False)
        tokenizer.load(tokenizer_path)
    else:
        tokenizer = KinTokenizer()

    if tokenizer.vocab_size is None or not tokenizer.merged_tokens:
        raise ValueError(
            "Tokenizer has no vocabulary/merges. Provide a valid tokenizer via "
            "tokenizer_path or ensure data/checkpoint/kin_tokenizer.json exists."
        )

    if not 0 <= eval_ratio < 1:
        raise ValueError("eval_ratio must be in the range [0, 1)")

    if not os.path.exists(destination_dir) or not os.path.isdir(destination_dir):
        raise ValueError("The destination folder/directory should be valid and exist")

    # ── Phase 1: Load text ──
    phase_start = time.time()
    if is_text_file_path_url:
        try:
            text = requests.get(text_file_path, timeout=120).text
        except Exception as e:
            raise ValueError(f"Unable to read the text.\n{str(e)}")
    else:
        if not os.path.exists(text_file_path) or not os.path.isfile(text_file_path):
            raise ValueError("The text file path should be valid like folder/text.txt")
        text = _read_text_efficient(text_file_path)
    print(f"Text loaded: {len(text):,} characters ({time.time() - phase_start:.1f}s)")

    if is_lowercase_text is None:
        is_lowercase_text = bool(tokenizer.normalization.get("lowercase", False))

    # ── Phase 2: Preprocess ──
    phase_start = time.time()
    text = preprocess_text(text, is_lowercase_text)
    print(f"Preprocessing done: {len(text):,} characters ({time.time() - phase_start:.1f}s)")

    if step_size is None:
        step_size = max(1, sequence_length // 2)

    # ── Phase 3: Train / Eval paragraph split ──
    if eval_ratio > 0:
        phase_start = time.time()
        train_text, eval_text = _split_paragraphs_train_eval(text, eval_ratio, eval_seed)
        del text
        gc.collect()
        train_chars = len(train_text)
        eval_chars = len(eval_text)
        actual_ratio = eval_chars / (train_chars + eval_chars) if (train_chars + eval_chars) > 0 else 0
        print(f"Paragraph-level split done ({time.time() - phase_start:.1f}s):")
        print(f"  Train: {train_chars:,} chars ({100 - actual_ratio * 100:.1f}%)")
        print(f"  Eval:  {eval_chars:,} chars ({actual_ratio * 100:.1f}%)")
        print(f"  Seed:  {eval_seed}")
    else:
        train_text = text
        eval_text = ""
        del text

    print("Sample text (train):\n\n", train_text[:sequence_length])
    print("\n" + "=" * 100 + "\n")

    # Adjust worker count based on text size
    ref_len = len(train_text)
    if ref_len < sequence_length + 2:
        nbr_processes = 1
    elif ref_len <= ((sequence_length + 1) * nbr_processes):
        nbr_processes = max(1, ref_len // (sequence_length + 1))
    else:
        nbr_processes = max(1, nbr_processes)

    # Detect Rust encoding path (shared for both splits)
    _use_rust_raw = (
        _rust_encode_text_raw is not None
        and getattr(tokenizer, "space_strategy", "metaspace") == "metaspace"
    )
    _use_rust_raw_batch = _use_rust_raw and _rust_encode_text_raw_batch is not None
    if _use_rust_raw_batch:
        print("Encoding path: Rust raw-bytes batch (single merge-map build)")
    elif _use_rust_raw:
        print("Encoding path: Rust raw-bytes (zero Python-int overhead)")
    else:
        print("Encoding path: chunked tokenizer.encode()")

    # ── Phase 4: Encode & write TRAIN sequences ──
    print(f"\n--- Encoding TRAIN split ({len(train_text):,} chars) ---")
    train_text = tokenizer._normalize_text_for_encoding(train_text, is_lowercase_text, _skip_nfc=True)
    train_tokens, train_enc_time = _encode_text_to_tokens(train_text, tokenizer, nbr_processes, label="train")
    del train_text
    gc.collect()
    print(f"Train encoding: {len(train_tokens):,} tokens, "
          f"{train_tokens.nbytes / (1024**3):.2f} GB ({train_enc_time:.1f}s)")

    train_file = os.path.join(destination_dir, "train_sequences.npy")
    print(f"Creating train sequences (seq_len={sequence_length}, step={step_size})...")
    train_seqs, _ = _write_sequences_npy(train_tokens, sequence_length, step_size, train_file, label="train")
    del train_tokens
    gc.collect()

    if train_seqs == 0:
        raise ValueError(f"No train sequences generated — need at least {sequence_length + 1} tokens")

    # ── Phase 5: Encode & write EVAL sequences ──
    eval_seqs = 0
    if eval_text:
        print(f"\n--- Encoding EVAL split ({len(eval_text):,} chars) ---")
        eval_text = tokenizer._normalize_text_for_encoding(eval_text, is_lowercase_text, _skip_nfc=True)
        eval_tokens, eval_enc_time = _encode_text_to_tokens(eval_text, tokenizer, nbr_processes, label="eval")
        del eval_text
        gc.collect()
        print(f"Eval encoding: {len(eval_tokens):,} tokens, "
              f"{eval_tokens.nbytes / (1024**3):.2f} GB ({eval_enc_time:.1f}s)")

        eval_file = os.path.join(destination_dir, "eval_sequences.npy")
        print(f"Creating eval sequences (seq_len={sequence_length}, step={step_size})...")
        eval_seqs, _ = _write_sequences_npy(eval_tokens, sequence_length, step_size, eval_file, label="eval")
        del eval_tokens
        gc.collect()

        if eval_seqs == 0:
            print("WARNING: Eval split too small to form sequences — only train_sequences.npy produced")
    else:
        print("\nNo eval split (eval_ratio=0)")

    # ── Summary ──
    total_time = time.time() - start_time
    training_end_time = datetime.datetime.now(datetime.timezone.utc)
    formatted_time = _format_duration(total_time)

    print(f"\n{'=' * 60}")
    print(f"  Dataset Creation Summary")
    print(f"{'=' * 60}")
    print(f"  Train sequences: {train_seqs:,}")
    if eval_seqs > 0:
        print(f"  Eval sequences:  {eval_seqs:,}")
        print(f"  Eval ratio:      {eval_ratio:.1%} (seed={eval_seed})")
    print(f"  Sequence length: {sequence_length} (+1 target)")
    print(f"  Step size:       {step_size}")
    print(f"  Output dir:      {destination_dir}")
    print(f"  Started:         {training_start_time.strftime('%d-%m-%Y %H:%M:%S')} UTC")
    print(f"  Finished:        {training_end_time.strftime('%d-%m-%Y %H:%M:%S')} UTC")
    print(f"  Total time:      {formatted_time}")
    print(f"{'=' * 60}\n")
