"""
Real-time streaming comparison: nanotok vs tiktoken.

Encodes and decodes text chunk-by-chunk, printing side-by-side timing
as each chunk completes — so you can watch the race live.

Usage:
  python benchmarks/stream_compare.py
  python benchmarks/stream_compare.py --encoding o200k_base
  python benchmarks/stream_compare.py --chunk-size 200
"""

from __future__ import annotations

import argparse
import sys
import time

TEXT = """\
The emergence of large language models has fundamentally transformed the landscape \
of artificial intelligence research and application development. Models like GPT-4, \
Claude, and LLaMA have demonstrated remarkable capabilities in natural language \
understanding, code generation, mathematical reasoning, and creative writing.

However, the computational cost of running these models remains a significant \
bottleneck. Tokenization — the process of converting raw text into integer sequences \
that models can process — is often overlooked as a performance-critical component. \
While attention mechanisms and matrix multiplications receive most of the optimization \
effort, tokenization can account for a non-trivial fraction of end-to-end latency, \
especially in streaming and real-time applications.

def quicksort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return quicksort(left) + middle + quicksort(right)

SELECT u.id, u.name, u.email, COUNT(o.id) AS total_orders,
       SUM(o.amount) AS total_spent, AVG(o.amount) AS avg_order
FROM users u
JOIN orders o ON u.id = o.user_id
WHERE o.created_at >= '2024-01-01'
  AND o.status IN ('completed', 'shipped')
GROUP BY u.id, u.name, u.email
HAVING COUNT(o.id) > 3
ORDER BY total_spent DESC
LIMIT 50;

The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs. \
How vexingly quick daft zebras jump! The five boxing wizards jump quickly. \
こんにちは世界！今日は素晴らしい日です。Привет мир! 你好世界 🌍🚀✨ \
€1,234.56 | £999.99 | ¥10,000 | ₹5,000 | 안녕하세요

{"model": "gpt-4o", "messages": [{"role": "system", "content": "You are a helpful assistant."}, \
{"role": "user", "content": "Explain quantum computing in simple terms."}], \
"temperature": 0.7, "max_tokens": 1024, "stream": true}
"""

RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
CYAN = "\033[36m"
YELLOW = "\033[33m"
RED = "\033[31m"
MAGENTA = "\033[35m"


def bar(value: float, max_val: float, width: int = 30, color: str = GREEN) -> str:
    if max_val <= 0:
        return " " * width
    filled = int(value / max_val * width)
    filled = min(filled, width)
    return f"{color}{'█' * filled}{DIM}{'░' * (width - filled)}{RESET}"


def stream_encode(text: str, encoding: str, chunk_size: int):
    import tiktoken
    from nanotok import Tokenizer

    tok_tt = tiktoken.get_encoding(encoding)
    tok_nt = Tokenizer.from_tiktoken(encoding)

    chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

    print(f"\n{BOLD}{'='*80}{RESET}")
    print(f"{BOLD}  STREAMING ENCODE — {encoding}  ({len(text)} chars, {len(chunks)} chunks of ~{chunk_size}){RESET}")
    print(f"{BOLD}{'='*80}{RESET}\n")

    header = f"  {'Chunk':>5s}  {'nanotok':>10s}  {'tiktoken':>10s}  {'Speedup':>8s}  {'Parity':>7s}  {'Tokens':>7s}  nanotok vs tiktoken"
    print(f"{DIM}{header}{RESET}")
    print(f"  {'─'*90}")

    nt_total_us, tt_total_us = 0.0, 0.0
    nt_total_tokens, tt_total_tokens = 0, 0
    mismatches = 0

    for i, chunk in enumerate(chunks):
        # nanotok
        t0 = time.perf_counter_ns()
        nt_ids = tok_nt.encode(chunk)
        nt_us = (time.perf_counter_ns() - t0) / 1000

        # tiktoken
        t0 = time.perf_counter_ns()
        tt_ids = tok_tt.encode(chunk)
        tt_us = (time.perf_counter_ns() - t0) / 1000

        match = nt_ids == tt_ids
        if not match:
            mismatches += 1

        speedup = tt_us / nt_us if nt_us > 0 else 0
        nt_total_us += nt_us
        tt_total_us += tt_us
        nt_total_tokens += len(nt_ids)
        tt_total_tokens += len(tt_ids)

        max_us = max(nt_us, tt_us)
        nt_bar = bar(nt_us, max_us, 15, CYAN)
        tt_bar = bar(tt_us, max_us, 15, YELLOW)

        parity_str = f"{GREEN}  ✓{RESET}" if match else f"{RED}  ✗{RESET}"
        spd_color = GREEN if speedup >= 1 else RED
        spd_str = f"{spd_color}{speedup:>6.1f}x{RESET}"

        print(
            f"  {i+1:>5d}  {nt_us:>8.0f}µs  {tt_us:>8.0f}µs  {spd_str}  {parity_str}  {len(nt_ids):>7d}  {nt_bar} {tt_bar}",
            flush=True,
        )
        time.sleep(0.03)

    total_speedup = tt_total_us / nt_total_us if nt_total_us > 0 else 0
    print(f"  {'─'*90}")
    print(f"  {'TOTAL':>5s}  {nt_total_us:>8.0f}µs  {tt_total_us:>8.0f}µs  {GREEN}{total_speedup:>6.1f}x{RESET}  {'':>7s}  {nt_total_tokens:>7d}")
    print(f"\n  {BOLD}Encode: nanotok is {total_speedup:.1f}x faster than tiktoken{RESET}")
    print(f"  Parity: {len(chunks) - mismatches}/{len(chunks)} chunks match ({(len(chunks) - mismatches) / len(chunks) * 100:.1f}%)")

    return tok_nt, tok_tt, nt_total_tokens


def stream_decode(tok_nt, tok_tt, text: str, encoding: str, chunk_size: int):
    chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

    all_nt_ids = [tok_nt.encode(c) for c in chunks]
    all_tt_ids = [tok_tt.encode(c) for c in chunks]

    print(f"\n{BOLD}{'='*80}{RESET}")
    print(f"{BOLD}  STREAMING DECODE — {encoding}{RESET}")
    print(f"{BOLD}{'='*80}{RESET}\n")

    header = f"  {'Chunk':>5s}  {'nanotok':>10s}  {'tiktoken':>10s}  {'Speedup':>8s}  {'Parity':>7s}  nanotok vs tiktoken"
    print(f"{DIM}{header}{RESET}")
    print(f"  {'─'*80}")

    nt_total_us, tt_total_us = 0.0, 0.0
    mismatches = 0

    for i, (nt_ids, tt_ids) in enumerate(zip(all_nt_ids, all_tt_ids)):
        t0 = time.perf_counter_ns()
        nt_text = tok_nt.decode(nt_ids)
        nt_us = (time.perf_counter_ns() - t0) / 1000

        t0 = time.perf_counter_ns()
        tt_text = tok_tt.decode(tt_ids)
        tt_us = (time.perf_counter_ns() - t0) / 1000

        match = nt_text == tt_text
        if not match:
            mismatches += 1

        speedup = tt_us / nt_us if nt_us > 0 else 0
        nt_total_us += nt_us
        tt_total_us += tt_us

        max_us = max(nt_us, tt_us)
        nt_bar = bar(nt_us, max_us, 15, CYAN)
        tt_bar = bar(tt_us, max_us, 15, YELLOW)

        parity_str = f"{GREEN}  ✓{RESET}" if match else f"{RED}  ✗{RESET}"
        spd_color = GREEN if speedup >= 1 else RED
        spd_str = f"{spd_color}{speedup:>6.1f}x{RESET}"

        print(
            f"  {i+1:>5d}  {nt_us:>8.0f}µs  {tt_us:>8.0f}µs  {spd_str}  {parity_str}  {nt_bar} {tt_bar}",
            flush=True,
        )
        time.sleep(0.03)

    total_speedup = tt_total_us / nt_total_us if nt_total_us > 0 else 0
    print(f"  {'─'*80}")
    print(f"  {'TOTAL':>5s}  {nt_total_us:>8.0f}µs  {tt_total_us:>8.0f}µs  {GREEN}{total_speedup:>6.1f}x{RESET}")
    print(f"\n  {BOLD}Decode: nanotok is {total_speedup:.1f}x faster than tiktoken{RESET}")
    print(f"  Parity: {len(chunks) - mismatches}/{len(chunks)} chunks match ({(len(chunks) - mismatches) / len(chunks) * 100:.1f}%)")


def stream_token_by_token_decode(tok_nt, tok_tt, text: str):
    ids = tok_tt.encode(text)
    n = len(ids)

    print(f"\n{BOLD}{'='*80}{RESET}")
    print(f"{BOLD}  TOKEN-BY-TOKEN DECODE RACE — simulating LLM streaming output ({n} tokens){RESET}")
    print(f"{BOLD}{'='*80}{RESET}\n")

    print(f"  {CYAN}nanotok:{RESET}  ", end="", flush=True)
    nt_total_ns = 0
    for i in range(1, n + 1):
        t0 = time.perf_counter_ns()
        try:
            tok_nt.decode(ids[:i])
        except UnicodeDecodeError:
            pass
        nt_total_ns += time.perf_counter_ns() - t0

    partial = tok_nt.decode(ids)
    preview = partial[:80].replace("\n", "\\n")
    print(f"{preview}... ({nt_total_ns / 1e6:.2f} ms total)")

    print(f"  {YELLOW}tiktoken:{RESET} ", end="", flush=True)
    tt_total_ns = 0
    for i in range(1, n + 1):
        t0 = time.perf_counter_ns()
        try:
            tok_tt.decode(ids[:i])
        except UnicodeDecodeError:
            pass
        tt_total_ns += time.perf_counter_ns() - t0

    partial = tok_tt.decode(ids)
    preview = partial[:80].replace("\n", "\\n")
    print(f"{preview}... ({tt_total_ns / 1e6:.2f} ms total)")

    speedup = tt_total_ns / nt_total_ns if nt_total_ns > 0 else 0
    print(f"\n  {BOLD}Token-by-token decode: nanotok is {speedup:.1f}x faster{RESET}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--encoding", default="cl100k_base", help="tiktoken encoding name")
    ap.add_argument("--chunk-size", type=int, default=150, help="chars per streaming chunk")
    args = ap.parse_args()

    print(f"\n{BOLD}{MAGENTA}  nanotok vs tiktoken — Real-Time Streaming Comparison{RESET}")
    print(f"{DIM}  encoding: {args.encoding} | chunk size: {args.chunk_size} chars | text: {len(TEXT)} chars{RESET}")

    tok_nt, tok_tt, _ = stream_encode(TEXT, args.encoding, args.chunk_size)
    stream_decode(tok_nt, tok_tt, TEXT, args.encoding, args.chunk_size)
    stream_token_by_token_decode(tok_nt, tok_tt, TEXT)

    print()


if __name__ == "__main__":
    main()
