"""DAPT/TAPT corpus loading and tokenization.

Loads clean text from JSONL (sharded directory or single file),
tokenizes, and concatenates + chunks into fixed-length sequences
for masked language modeling.
"""

from os import cpu_count
from pathlib import Path

from datasets import Dataset, load_dataset
from transformers import PreTrainedTokenizerBase


def load_corpus(corpus_path: str, text_field: str = "text") -> Dataset:
    """Load JSONL file(s) into a HuggingFace Dataset.

    Args:
        corpus_path: Path to a single JSONL file or a directory of JSONL shards.
        text_field: Name of the JSON field containing the text.

    Returns:
        A Dataset with at least a `text_field` column.
    """
    p = Path(corpus_path)
    if p.is_file():
        ds = load_dataset("json", data_files=str(p), split="train")
    elif p.is_dir():
        ds = load_dataset("json", data_dir=str(p), split="train")
    else:
        raise FileNotFoundError(f"Corpus path not found: {corpus_path}")

    if text_field not in ds.column_names:
        raise ValueError(
            f"Text field '{text_field}' not found. "
            f"Available columns: {ds.column_names}"
        )
    return ds


def tokenize_and_chunk(
    dataset: Dataset,
    tokenizer: PreTrainedTokenizerBase,
    text_field: str = "text",
    max_seq_length: int = 2048,
    whole_word_mask: bool = False,
) -> Dataset:
    """Tokenize texts for MLM pre-training.

    Two modes:
    - **DAPT** (whole_word_mask=False): Concatenate all tokens and chunk into
      fixed-length sequences. Standard for pre-training on long documents —
      maximizes GPU utilization by avoiding padding.
    - **TAPT** (whole_word_mask=True): Tokenize each text individually with
      truncation, preserving word_ids for whole-word masking. The data
      collator handles dynamic padding per batch.

    Args:
        dataset: Dataset with a text column.
        tokenizer: Pretrained tokenizer.
        text_field: Column name containing raw text.
        max_seq_length: Length of each output sequence.
        whole_word_mask: If True, return per-sample word_ids instead
            of concatenate-and-chunk. Required for whole-word masking.

    Returns:
        Dataset with `input_ids` and `attention_mask` columns.
        If whole_word_mask=True, also includes `word_ids`.
    """
    nproc = cpu_count() or 1

    if whole_word_mask:
        # Tokenize each paragraph individually, extracting word_ids for
        # whole-word masking. We use word_ids() (from the fast tokenizer)
        # instead of offset_mapping because BPE tokenizers absorb leading
        # spaces into tokens, making all offsets contiguous — the upstream
        # offset_mapping-based boundary detection treats entire sequences
        # as one "word".
        def tokenize_fn(examples: dict) -> dict:
            encodings = tokenizer(
                examples[text_field],
                truncation=True,
                max_length=max_seq_length,
            )
            batch_word_ids = []
            for i in range(len(examples[text_field])):
                wids = encodings.word_ids(i)
                # None (special tokens) → -1 for Arrow integer storage
                batch_word_ids.append([w if w is not None else -1 for w in wids])
            return {
                "input_ids": encodings["input_ids"],
                "attention_mask": encodings["attention_mask"],
                "word_ids": batch_word_ids,
            }

        return dataset.map(
            tokenize_fn,
            batched=True,
            batch_size=100,
            num_proc=nproc,
            remove_columns=dataset.column_names,
            desc="Tokenizing (whole-word mask)",
        )

    # DAPT mode: concatenate-then-chunk
    def tokenize_fn(examples: dict) -> dict:
        return tokenizer(
            examples[text_field],
            truncation=False,
            add_special_tokens=False,
        )

    tokenized = dataset.map(
        tokenize_fn,
        batched=True,
        batch_size=100,
        num_proc=nproc,
        remove_columns=dataset.column_names,
        desc="Tokenizing",
    )

    # Concatenate-then-chunk in parallel batches. Each batch of documents
    # has its tokens concatenated and sliced into fixed-length sequences.
    # Only the tail of each batch is dropped (not each document), so token
    # waste is negligible (~nproc * max_seq_length tokens total).
    def chunk_fn(examples: dict) -> dict:
        all_ids: list[int] = []
        all_mask: list[int] = []
        for ids, mask in zip(examples["input_ids"], examples["attention_mask"]):
            all_ids.extend(ids)
            all_mask.extend(mask)
        n = (len(all_ids) // max_seq_length) * max_seq_length
        return {
            "input_ids": [all_ids[i : i + max_seq_length] for i in range(0, n, max_seq_length)],
            "attention_mask": [all_mask[i : i + max_seq_length] for i in range(0, n, max_seq_length)],
        }

    return tokenized.map(
        chunk_fn, batched=True, batch_size=500, num_proc=nproc,
        remove_columns=tokenized.column_names, desc="Chunking",
    )
