#!/usr/bin/env python3
"""
Tokenizer Demo — Lesson 01: How LLMs Work

Prints a colourised token breakdown for any input string.

Usage:
    python demo/tokenizer_demo.py "Your text here"
    python demo/tokenizer_demo.py "def hello(): return 'world'"
    python demo/tokenizer_demo.py  # uses a built-in example
"""

import sys
import tiktoken

# ---------------------------------------------------------------------------
# ANSI colour codes for terminal output
# ---------------------------------------------------------------------------
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"

# Eight distinct background colours for token highlighting
TOKEN_COLOURS = [
    "\033[41m",   # red bg
    "\033[44m",   # blue bg
    "\033[42m",   # green bg
    "\033[43m",   # yellow bg
    "\033[45m",   # magenta bg
    "\033[46m",   # cyan bg
    "\033[101m",  # bright red bg
    "\033[104m",  # bright blue bg
]

# Foreground colours paired with each background for readability
TOKEN_FG = [
    "\033[97m",   # white text on red
    "\033[97m",   # white text on blue
    "\033[30m",   # black text on green
    "\033[30m",   # black text on yellow
    "\033[97m",   # white text on magenta
    "\033[30m",   # black text on cyan
    "\033[97m",   # white text on bright red
    "\033[97m",   # white text on bright blue
]


def colour_token(text: str, index: int) -> str:
    """Wrap token text in ANSI colour codes, cycling through the palette."""
    colour_index = index % len(TOKEN_COLOURS)
    bg = TOKEN_COLOURS[colour_index]
    fg = TOKEN_FG[colour_index]
    # Replace spaces with a visible marker so whitespace tokens are visible
    display_text = text.replace(" ", "·").replace("\n", "↵").replace("\t", "→")
    return f"{bg}{fg}{display_text}{RESET}"


def tokenize_and_display(text: str, encoding_name: str = "cl100k_base") -> None:
    """
    Tokenise the input text and print a detailed, colourised breakdown.

    Shows:
    - Colourised inline token display
    - Token ID table
    - Summary statistics
    - Cost estimate
    """
    enc = tiktoken.get_encoding(encoding_name)
    token_ids = enc.encode(text)
    tokens = [enc.decode([tid]) for tid in token_ids]

    # -----------------------------------------------------------------------
    # Header
    # -----------------------------------------------------------------------
    print()
    print(f"{BOLD}Tokenizer Demo — {encoding_name}{RESET}")
    print("─" * 60)
    print(f"{BOLD}Input:{RESET}")
    print(f"  {text!r}")
    print()

    # -----------------------------------------------------------------------
    # Colourised inline view
    # -----------------------------------------------------------------------
    print(f"{BOLD}Tokens (colourised):{RESET}")
    print("  ", end="")
    for i, token_text in enumerate(tokens):
        print(colour_token(token_text, i), end="")
    print()
    print()

    # -----------------------------------------------------------------------
    # Token ID table
    # -----------------------------------------------------------------------
    print(f"{BOLD}{'#':<5} {'Token ID':>10}  {'Token text':<30}{RESET}")
    print("─" * 50)
    for i, (tid, token_text) in enumerate(zip(token_ids, tokens)):
        # Show spaces, newlines, tabs visually
        display = token_text.replace(" ", "·").replace("\n", "↵").replace("\t", "→")
        colour_index = i % len(TOKEN_COLOURS)
        bg = TOKEN_COLOURS[colour_index]
        fg = TOKEN_FG[colour_index]
        coloured_display = f"{bg}{fg}{display:<20}{RESET}"
        print(f"  {i:<3}  {tid:>10}  {coloured_display}")
    print()

    # -----------------------------------------------------------------------
    # Summary statistics
    # -----------------------------------------------------------------------
    n_tokens = len(token_ids)
    n_chars = len(text)
    n_words = len(text.split())
    chars_per_token = n_chars / n_tokens if n_tokens else 0
    tokens_per_word = n_tokens / n_words if n_words else 0

    print(f"{BOLD}Summary:{RESET}")
    print(f"  Characters : {n_chars}")
    print(f"  Words      : {n_words}")
    print(f"  Tokens     : {n_tokens}")
    print(f"  Chars/token: {chars_per_token:.2f}")
    print(f"  Tokens/word: {tokens_per_word:.2f}  (rule of thumb: ~1.3 for English prose)")
    print()

    # -----------------------------------------------------------------------
    # Quick cost estimate
    # -----------------------------------------------------------------------
    # Prices last verified 2024-10-22 — check https://www.anthropic.com/pricing for updates
    # Pricing reference (USD per million tokens, approximate)
    PRICING = {
        "claude-3-5-sonnet-20241022": {"input": 3.00, "output": 15.00},
        "claude-3-5-haiku-20241022":  {"input": 0.80, "output": 4.00},
    }

    print(f"{BOLD}Cost estimate (input tokens only, per 1000 calls):{RESET}")
    print(f"  {'Model':<40} {'$/1k calls':>12}")
    print("  " + "─" * 54)
    for model, prices in PRICING.items():
        cost_per_1k = (n_tokens / 1_000_000) * prices["input"] * 1000
        print(f"  {model:<40} ${cost_per_1k:>10.4f}")
    print()


def show_comparison(strings: list[str], encoding_name: str = "cl100k_base") -> None:
    """Show a side-by-side token count comparison for multiple strings."""
    enc = tiktoken.get_encoding(encoding_name)

    print(f"{BOLD}Token Count Comparison:{RESET}")
    print(f"  {'String':<45} {'Tokens':>7}")
    print("  " + "─" * 54)
    for s in strings:
        count = len(enc.encode(s))
        display = s if len(s) <= 43 else s[:40] + "..."
        print(f"  {display!r:<45} {count:>7}")
    print()


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
DEFAULT_TEXT = (
    "The quick brown fox jumps over the lazy dog. "
    "In 2024, large language models changed how we write software."
)

COMPARISON_EXAMPLES = [
    "Hello",
    " Hello",           # leading space — different token!
    "hello",            # lowercase
    "Hello, world!",
    "The quick brown fox",
    "def calculate_fibonacci(n: int) -> int:",
    "Привет, мир!",     # Russian
    "2024",
    "20240101",
    "   ",              # 3 spaces
]


def main() -> None:
    # If a string is provided as a command-line argument, use it; otherwise use the default
    if len(sys.argv) > 1:
        input_text = " ".join(sys.argv[1:])
    else:
        print(f"{DIM}No input provided. Using built-in example. Usage: python demo/tokenizer_demo.py \"your text\"{RESET}")
        input_text = DEFAULT_TEXT

    tokenize_and_display(input_text)

    # Always show the comparison table — it's educational
    print("─" * 60)
    show_comparison(COMPARISON_EXAMPLES)


if __name__ == "__main__":
    main()