#!/usr/bin/env python3
"""
Tiktoken-style benchmark comparing SARFTokenizer vs tiktoken vs HuggingFace.

Measures throughput in MB/s with proper thread isolation using multiprocessing.

Usage:
    python benchmark_tiktoken_style.py --samples 1000000 --threads 1 2 4 8
"""

import os
import sys
import time
import argparse
from pathlib import Path
from typing import List, Tuple
from multiprocessing import Process, Queue, cpu_count

import pyarrow.parquet as pq

# Add parent to path
sys.path.insert(0, str(Path(__file__).parent))

# Configuration
DATA_DIR = "/root/.cache/deeplatent/base_data/"
HF_TOKENIZER_PATH = os.path.expanduser("~/.cache/deeplatent/tokenizers/SARFTokenizer")
DEFAULT_THREADS = [2**i for i in range(8) if 2**i <= cpu_count()]


def format_byte_size(num_bytes: float) -> Tuple[str, str]:
    """Convert bytes to human-readable format."""
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if num_bytes < 1024:
            return f"{num_bytes:.2f} {unit}", unit
        num_bytes /= 1024
    return f"{num_bytes:.2f} PB", "PB"


def load_samples(data_dir: str, num_samples: int) -> Tuple[List[str], int]:
    """Load samples from parquet files."""
    import re
    AR_DETECT = re.compile(r'[\u0600-\u06FF]')

    parquet_files = sorted(Path(data_dir).glob("shard_*.parquet"))
    if not parquet_files:
        raise FileNotFoundError(f"No parquet files found in {data_dir}")

    samples = []
    target = num_samples

    for pq_file in parquet_files:
        if len(samples) >= target:
            break

        table = pq.read_table(pq_file, columns=["text"])
        texts = table.column("text").to_pylist()

        for text in texts:
            if len(samples) >= target:
                break
            if text and isinstance(text, str):
                samples.append(text)

    total_bytes = sum(len(t.encode('utf-8')) for t in samples)
    return samples, total_bytes


def benchmark_sarf(documents: List[str], num_threads: int, result_queue: Queue):
    """Benchmark SARFTokenizer."""
    from deeplatent import SARFTokenizer

    os.environ["RAYON_NUM_THREADS"] = str(num_threads)

    tok = SARFTokenizer.from_pretrained(HF_TOKENIZER_PATH)
    num_bytes = sum(len(d.encode('utf-8')) for d in documents)

    # Warmup
    tok.encode(documents[0])

    # Benchmark
    start = time.perf_counter_ns()
    if hasattr(tok, 'encode_batch'):
        tok.encode_batch(documents)
    else:
        for d in documents:
            tok.encode(d)
    end = time.perf_counter_ns()

    elapsed_ns = end - start
    bytes_per_sec = num_bytes / elapsed_ns * 1e9
    texts_per_sec = len(documents) / elapsed_ns * 1e9

    result_queue.put(("SARFTokenizer", bytes_per_sec, texts_per_sec))


def benchmark_tiktoken(documents: List[str], num_threads: int, encoding: str, result_queue: Queue):
    """Benchmark tiktoken."""
    import tiktoken

    os.environ["RAYON_NUM_THREADS"] = str(num_threads)

    enc = tiktoken.get_encoding(encoding)
    num_bytes = sum(len(d.encode('utf-8')) for d in documents)

    # Warmup
    enc.encode(documents[0])

    # Benchmark
    start = time.perf_counter_ns()
    enc.encode_ordinary_batch(documents, num_threads=num_threads)
    end = time.perf_counter_ns()

    elapsed_ns = end - start
    bytes_per_sec = num_bytes / elapsed_ns * 1e9
    texts_per_sec = len(documents) / elapsed_ns * 1e9

    result_queue.put((f"tiktoken ({encoding})", bytes_per_sec, texts_per_sec))


def benchmark_hf_tokenizers(documents: List[str], num_threads: int, result_queue: Queue):
    """Benchmark HuggingFace tokenizers."""
    from tokenizers import Tokenizer

    os.environ["RAYON_NUM_THREADS"] = str(num_threads)

    # Load the SARFTokenizer's underlying HF tokenizer
    tokenizer_path = os.path.join(HF_TOKENIZER_PATH, "tokenizer.json")
    tok = Tokenizer.from_file(tokenizer_path)
    num_bytes = sum(len(d.encode('utf-8')) for d in documents)

    # Warmup
    tok.encode(documents[0])

    # Benchmark
    start = time.perf_counter_ns()
    tok.encode_batch_fast(documents)
    end = time.perf_counter_ns()

    elapsed_ns = end - start
    bytes_per_sec = num_bytes / elapsed_ns * 1e9
    texts_per_sec = len(documents) / elapsed_ns * 1e9

    result_queue.put(("HF tokenizers", bytes_per_sec, texts_per_sec))


def run_benchmark(documents: List[str], num_threads: int, num_bytes: int):
    """Run benchmarks for all tokenizers with given thread count."""
    readable_size, _ = format_byte_size(num_bytes)
    avg_len = sum(len(d) for d in documents) / len(documents)

    print(f"\n{'='*70}")
    print(f"Threads: {num_threads}, Data: {readable_size}, Documents: {len(documents):,}, Avg Length: {avg_len:.0f}")
    print(f"{'='*70}")

    results = []

    # SARFTokenizer
    q = Queue()
    p = Process(target=benchmark_sarf, args=(documents, num_threads, q))
    p.start()
    p.join()
    if not q.empty():
        name, bps, tps = q.get()
        readable, _ = format_byte_size(bps)
        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
        results.append((name, bps, tps))

    # tiktoken o200k_base
    q = Queue()
    p = Process(target=benchmark_tiktoken, args=(documents, num_threads, "o200k_base", q))
    p.start()
    p.join()
    if not q.empty():
        name, bps, tps = q.get()
        readable, _ = format_byte_size(bps)
        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
        results.append((name, bps, tps))

    # tiktoken cl100k_base
    q = Queue()
    p = Process(target=benchmark_tiktoken, args=(documents, num_threads, "cl100k_base", q))
    p.start()
    p.join()
    if not q.empty():
        name, bps, tps = q.get()
        readable, _ = format_byte_size(bps)
        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
        results.append((name, bps, tps))

    # HuggingFace tokenizers
    q = Queue()
    p = Process(target=benchmark_hf_tokenizers, args=(documents, num_threads, q))
    p.start()
    p.join()
    if not q.empty():
        name, bps, tps = q.get()
        readable, _ = format_byte_size(bps)
        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
        results.append((name, bps, tps))

    return results


def main():
    parser = argparse.ArgumentParser(description="Tiktoken-style tokenizer benchmark")
    parser.add_argument("--samples", type=int, default=10000, help="Number of samples")
    parser.add_argument("--threads", type=int, nargs="+", default=DEFAULT_THREADS, help="Thread counts")
    parser.add_argument("--data-dir", type=str, default=DATA_DIR, help="Data directory")
    args = parser.parse_args()

    print("=" * 70)
    print("TIKTOKEN-STYLE TOKENIZER BENCHMARK")
    print("=" * 70)
    print(f"CPU count: {cpu_count()}")
    print(f"Samples: {args.samples:,}")
    print(f"Threads: {args.threads}")

    # Load data
    print("\nLoading data...")
    documents, total_bytes = load_samples(args.data_dir, args.samples)
    readable_size, _ = format_byte_size(total_bytes)
    print(f"Loaded {len(documents):,} documents ({readable_size})")

    # Run benchmarks
    all_results = {}
    for num_threads in args.threads:
        results = run_benchmark(documents, num_threads, total_bytes)
        all_results[num_threads] = results

    # Summary table
    print("\n" + "=" * 100)
    print("SUMMARY TABLE (MB/s)")
    print("=" * 100)

    # Header
    header = f"{'Tokenizer':<25}"
    for t in args.threads:
        header += f"{t}T".rjust(15)
    print(header)
    print("-" * 100)

    # Collect by tokenizer name
    tokenizers = {}
    for threads, results in all_results.items():
        for name, bps, tps in results:
            if name not in tokenizers:
                tokenizers[name] = {}
            tokenizers[name][threads] = bps / 1024 / 1024  # Convert to MB/s

    # Print rows
    for name, thread_results in tokenizers.items():
        row = f"{name:<25}"
        for t in args.threads:
            if t in thread_results:
                row += f"{thread_results[t]:>14.2f}"
            else:
                row += "N/A".rjust(15)
        print(row)

    print("=" * 100)


if __name__ == "__main__":
    main()