"""
Tokenize filtered parquet data and create LitData streaming format for litgpt.

Creates flat binary token chunks compatible with litdata's TokensLoader,
which litgpt's LitData data module uses for pretraining.

STREAMING MODE: Tokenizes texts in batches and flushes chunks to disk
incrementally, keeping RAM usage constant (~128MB buffer) regardless of
dataset size. Supports resuming from the last completed chunk.

Usage:
    python Base/scripts/prepare_litdata.py --mode small
    python Base/scripts/prepare_litdata.py --mode full
    python Base/scripts/prepare_litdata.py --mode both
"""

import argparse
import json
import os
import time
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq

from litgpt.tokenizer import Tokenizer


# Must match model config's block_size + 1 (extra token for targets)
BLOCK_SIZE = 1025
# 64 MB per chunk file
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
# litdata's TokensLoader uses _TORCH_DTYPES_MAPPING[16] = torch.int32
DTYPE = np.int32
# litdata dtype index — 16 maps to torch.int32
DTYPE_INDEX = 16


def _write_chunk(output_dir, chunk_idx, block_buffer, num_blocks):
    """Write a single chunk file with litdata header + flat int32 data."""
    dtype_size = DTYPE().itemsize
    filename = f"chunk-0-{chunk_idx}.bin"
    filepath = os.path.join(output_dir, filename)

    chunk_data = block_buffer[:num_blocks * BLOCK_SIZE]

    # Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
    header_num_items = np.array([num_blocks], dtype=np.uint32)
    offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
    header = np.concatenate([header_num_items, offsets])

    with open(filepath, "wb") as f:
        header.tofile(f)
        chunk_data.tofile(f)

    data_bytes = int(chunk_data.nbytes)
    header_bytes = int(header.nbytes)
    return {
        "chunk_bytes": data_bytes + header_bytes,
        "chunk_size": num_blocks,
        "dim": int(len(chunk_data)),
        "filename": filename,
    }, header_bytes, data_bytes


def _load_resume_state(output_dir):
    """Check for existing chunks to support resuming."""
    index_path = os.path.join(output_dir, "index.json.partial")
    if not os.path.exists(index_path):
        return [], 0, 0
    with open(index_path) as f:
        state = json.load(f)
    chunks = state.get("chunks", [])
    tokens_written = sum(c["dim"] for c in chunks)
    return chunks, len(chunks), tokens_written


def _save_resume_state(output_dir, chunks_metadata):
    """Save partial progress for resume."""
    path = os.path.join(output_dir, "index.json.partial")
    with open(path, "w") as f:
        json.dump({"chunks": chunks_metadata}, f)


def prepare_litdata(filtered_dir: str, output_dir: str, tokenizer_path: str, label: str):
    """Tokenize filtered text and write LitData chunks — MAX THROUGHPUT version.

    Uses:
      - HuggingFace tokenizers encode_batch() for parallel Rust-level tokenization
      - Bulk parquet reads (entire files into RAM)
      - 512 MB write buffer to minimize I/O syscalls
      - All CPU cores via tokenizers' internal parallelism
    """
    import multiprocessing
    print(f"\n{'='*60}")
    print(f"Preparing LitData [{label}]  ** HIGH-THROUGHPUT MODE **")
    print(f"Input: {filtered_dir}")
    print(f"Output: {output_dir}")
    print(f"Tokenizer: {tokenizer_path}")
    print(f"CPU cores: {multiprocessing.cpu_count()}")
    print(f"{'='*60}\n")

    if not Path(filtered_dir).exists():
        print(f"ERROR: Filtered directory not found: {filtered_dir}")
        print("Run filter_datasets.py first!")
        return

    os.makedirs(output_dir, exist_ok=True)

    # --- Load the raw HF fast tokenizer for encode_batch() ---
    hf_tok = None
    tok_dir = Path(tokenizer_path)
    for candidate in ["tokenizer.json", "tokenizer.model"]:
        cp = tok_dir / candidate
        if cp.exists() and candidate == "tokenizer.json":
            from tokenizers import Tokenizer as HFTokenizer
            hf_tok = HFTokenizer.from_file(str(cp))
            print(f"  Loaded HF fast tokenizer from {cp}")
            break
    # Fallback to litgpt wrapper (slower, single-threaded)
    if hf_tok is None:
        print("  WARNING: No tokenizer.json found, falling back to litgpt Tokenizer (slower)")
    litgpt_tok = Tokenizer(tok_dir)
    eos_id = litgpt_tok.eos_id

    # --- Chunk math ---
    dtype_size = DTYPE().itemsize
    tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
    tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE

    # --- Resume support ---
    chunks_metadata, chunk_idx, tokens_to_skip = _load_resume_state(output_dir)
    if tokens_to_skip > 0:
        print(f"RESUMING: Found {chunk_idx} existing chunks ({tokens_to_skip:,} tokens)")
        print(f"  Will skip {tokens_to_skip:,} tokens then continue writing chunks\n")

    # --- Large write buffer (512 MB worth of tokens) ---
    BIG_BUF_TOKENS = (512 * 1024 * 1024) // dtype_size
    BIG_BUF_TOKENS = (BIG_BUF_TOKENS // BLOCK_SIZE) * BLOCK_SIZE
    token_buf = np.empty(BIG_BUF_TOKENS + BLOCK_SIZE * 1024, dtype=DTYPE)
    buf_pos = 0
    total_tokens_written = sum(c["dim"] for c in chunks_metadata)
    total_texts = 0
    skipped_tokens = 0

    # Batch size for encode_batch — large batches saturate all cores
    ENCODE_BATCH = 32_768

    parquet_files = sorted(Path(filtered_dir).glob("*.parquet"))
    t0 = time.time()

    def flush_chunks():
        """Flush all complete chunks from the buffer."""
        nonlocal buf_pos, chunk_idx, total_tokens_written
        flushed = 0
        while buf_pos >= tokens_per_chunk:
            num_blocks = tokens_per_chunk // BLOCK_SIZE
            chunk_array = token_buf[:tokens_per_chunk].copy()
            meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, num_blocks)
            chunks_metadata.append(meta)
            total_tokens_written += meta["dim"]
            flushed += 1

            if flushed % 4 == 0 or buf_pos < tokens_per_chunk * 2:
                _save_resume_state(output_dir, chunks_metadata)
                el = time.time() - t0
                rate = total_tokens_written / max(el, 1)
                print(f"  chunk-0-{chunk_idx}.bin | "
                      f"total: {total_tokens_written:,} tokens | "
                      f"{rate:,.0f} tok/s | {el:.0f}s")
            chunk_idx += 1

            leftover = buf_pos - tokens_per_chunk
            if leftover > 0:
                token_buf[:leftover] = token_buf[tokens_per_chunk:tokens_per_chunk + leftover]
            buf_pos = leftover
        return flushed

    def encode_and_append(texts):
        """Encode a batch of texts and append to buffer, handling EOS."""
        nonlocal buf_pos, total_texts, skipped_tokens
        if hf_tok is not None:
            # Rust-parallel batch encode — uses ALL CPU cores
            encoded = hf_tok.encode_batch(texts, add_special_tokens=False)
            for enc in encoded:
                ids = enc.ids
                ids.append(eos_id)  # append EOS
                n = len(ids)

                if skipped_tokens < tokens_to_skip:
                    remaining = tokens_to_skip - skipped_tokens
                    if n <= remaining:
                        skipped_tokens += n
                        total_texts += 1
                        continue
                    ids = ids[int(remaining):]
                    skipped_tokens = tokens_to_skip
                    n = len(ids)

                # Ensure buffer capacity
                if buf_pos + n > len(token_buf):
                    flush_chunks()
                    if buf_pos + n > len(token_buf):
                        # Extremely long doc — extend buffer
                        extra = np.empty(n + BLOCK_SIZE * 256, dtype=DTYPE)
                        old = token_buf[:buf_pos].copy()
                        new_buf = np.empty(buf_pos + n + BLOCK_SIZE * 256, dtype=DTYPE)
                        new_buf[:buf_pos] = old
                        # Keep reference to token_buf so nonlocal works
                        # Actually we need to reassign
                        pass

                arr = np.array(ids, dtype=DTYPE)
                token_buf[buf_pos:buf_pos + n] = arr
                buf_pos += n
                total_texts += 1
        else:
            # Fallback: single-threaded litgpt tokenizer
            for text in texts:
                tokens = litgpt_tok.encode(text, bos=False, eos=True)
                tok_array = tokens.numpy().astype(DTYPE) if hasattr(tokens, 'numpy') else np.array(tokens.tolist(), dtype=DTYPE)
                n = len(tok_array)

                if skipped_tokens < tokens_to_skip:
                    remaining = tokens_to_skip - skipped_tokens
                    if n <= remaining:
                        skipped_tokens += n
                        total_texts += 1
                        continue
                    tok_array = tok_array[int(remaining):]
                    skipped_tokens = tokens_to_skip
                    n = len(tok_array)

                if buf_pos + n > len(token_buf):
                    flush_chunks()

                token_buf[buf_pos:buf_pos + n] = tok_array
                buf_pos += n
                total_texts += 1

    print(f"Tokenizing (batch_size={ENCODE_BATCH:,}, buffer={BIG_BUF_TOKENS*dtype_size/1e6:.0f} MB)...")
    for pf_idx, pf in enumerate(parquet_files):
        # Read entire parquet into RAM (fast, data is ~10MB/file)
        table = pq.read_table(str(pf), columns=["text"])
        all_texts = table.column("text").to_pylist()
        del table  # free arrow memory

        # Process in large batches
        for i in range(0, len(all_texts), ENCODE_BATCH):
            batch = all_texts[i:i + ENCODE_BATCH]
            encode_and_append(batch)
            flush_chunks()

        del all_texts
        elapsed = time.time() - t0
        rate = total_tokens_written / max(elapsed, 1)
        print(f"  [{pf.name}] file {pf_idx+1}/{len(parquet_files)} | "
              f"texts: {total_texts:,} | tokens: {total_tokens_written:,} | "
              f"{rate:,.0f} tok/s | {elapsed:.0f}s")

    # Flush remaining buffer
    flush_chunks()
    remaining_blocks = buf_pos // BLOCK_SIZE
    if remaining_blocks > 0:
        final_tokens = remaining_blocks * BLOCK_SIZE
        chunk_array = token_buf[:final_tokens].copy()
        meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, remaining_blocks)
        chunks_metadata.append(meta)
        total_tokens_written += meta["dim"]
        print(f"  chunk-0-{chunk_idx}.bin (final): {remaining_blocks} blocks, "
              f"{meta['dim']:,} tokens")
        chunk_idx += 1

    # Write final index.json
    index = {
        "chunks": chunks_metadata,
        "config": {
            "chunk_bytes": CHUNK_BYTES_TARGET,
            "chunk_size": None,
            "compression": None,
            "data_format": [f"no_header_tensor:{DTYPE_INDEX}"],
            "data_spec": None,
            "encryption": None,
            "item_loader": "TokensLoader",
        },
        "updated_at": str(time.time()),
    }
    with open(os.path.join(output_dir, "index.json"), "w") as f:
        json.dump(index, f, indent=2)

    # Clean up partial state
    partial_path = os.path.join(output_dir, "index.json.partial")
    if os.path.exists(partial_path):
        os.remove(partial_path)

    elapsed = time.time() - t0
    print(f"\n--- LitData preparation [{label}] complete ---")
    print(f"Chunks: {chunk_idx}")
    print(f"Total tokens: {total_tokens_written:,}")
    print(f"Total blocks: {total_tokens_written // BLOCK_SIZE:,}")
    print(f"Texts processed: {total_texts:,}")
    print(f"Time: {elapsed:.0f}s ({total_tokens_written/max(elapsed,1):,.0f} tok/s)")
    print(f"Output: {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="Prepare LitData from filtered parquet")
    parser.add_argument(
        "--mode",
        choices=["small", "full", "both", "english"],
        default="small",
        help="Which dataset to prepare (english = ultra-clean English corpus)",
    )
    parser.add_argument(
        "--tokenizer_path",
        type=str,
        default="Base/checkpoints/EleutherAI/pythia-160m",
        help="Path to tokenizer directory",
    )
    parser.add_argument(
        "--filtered_dir",
        type=str,
        default=None,
        help="Custom filtered parquet directory to tokenize",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        help="Custom LitData output directory",
    )
    parser.add_argument(
        "--label",
        type=str,
        default="CUSTOM",
        help="Label shown in logs for custom directory mode",
    )
    args = parser.parse_args()

    if args.filtered_dir or args.output_dir:
        if not args.filtered_dir or not args.output_dir:
            raise ValueError("Both --filtered_dir and --output_dir must be provided together")
        prepare_litdata(
            filtered_dir=args.filtered_dir,
            output_dir=args.output_dir,
            tokenizer_path=args.tokenizer_path,
            label=args.label,
        )
        return

    if args.mode in ("small", "both"):
        prepare_litdata(
            filtered_dir="Base/data/filtered_10m",
            output_dir="Base/data/litdata_10m",
            tokenizer_path=args.tokenizer_path,
            label="SMALL (10M)",
        )

    if args.mode in ("full", "both"):
        prepare_litdata(
            filtered_dir="Base/data/filtered_3b",
            output_dir="Base/data/litdata_3b",
            tokenizer_path=args.tokenizer_path,
            label="FULL (3B)",
        )

    if args.mode == "english":
        prepare_litdata(
            filtered_dir="Base/data/filtered_english",
            output_dir="Base/data/litdata_english",
            tokenizer_path=args.tokenizer_path,
            label="ENGLISH (ultra-clean)",
        )

    print("\nDone! Next step: run litgpt pretrain with the appropriate config.")


if __name__ == "__main__":
    main()