#!/usr/bin/env python3
"""Tokenization demo: count tokens accurately, compare models, budget context.

Shows why words != tokens, how to count with tiktoken, and how to budget a
context window across system prompt + history + retrieved chunks + output.
"""
from __future__ import annotations

import sys


def estimate_tokens_heuristic(text: str) -> int:
    """Rough heuristic: ~4 chars/token for English (INACCURATE — demo only)."""
    return len(text) // 4


def main() -> None:
    # --- 1. Words vs tokens ---
    sample = "The quick brown fox jumps over the lazy dog. " * 20  # 180 words
    words = len(sample.split())
    heuristic = estimate_tokens_heuristic(sample)
    print("=== Words vs Tokens ===")
    print(f"Text: {words} words, {len(sample)} chars")
    print(f"Naive char heuristic (~4 chars/token): {heuristic} tokens")
    print("NOTE: heuristic is model-agnostic and can be off by 25%+.\n")

    # --- 2. Accurate counting with tiktoken ---
    try:
        import tiktoken
    except ImportError:
        print("Install tiktoken for accurate counts: pip install tiktoken")
        print("Skipping accurate count demo.")
        tiktoken = None

    if tiktoken:
        print("=== Accurate Token Counts ===")
        # Same text, different models -> different token counts
        models = ["gpt-4o", "gpt-4", "gpt-3.5-turbo"]
        for model_name in models:
            try:
                enc = tiktoken.encoding_for_model(model_name)
                count = len(enc.encode(sample))
                print(f"  {model_name:20s}: {count} tokens")
            except KeyError:
                print(f"  {model_name:20s}: (unknown model — use cl100k_base)")
        print()

        # --- 3. English vs code vs non-English ---
        print("=== Token Density by Content Type ===")
        english = "The company reported strong quarterly earnings growth."
        code = "def fibonacci(n):\n    return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)"
        non_english = "这是一个关于人工智能的中文句子。"  # Chinese
        enc = tiktoken.get_encoding("cl100k_base")

        for label, text in [("English", english), ("Code", code), ("Chinese", non_english)]:
            print(f"  {label:10s}: {len(text):3d} chars -> {len(enc.encode(text)):3d} tokens "
                  f"({len(text) / len(enc.encode(text)):.1f} chars/token)")

        # --- 4. Context budgeting ---
        print("\n=== Context Window Budgeting (128k model) ===")
        MAX_CONTEXT = 128_000
        system_prompt = "You are a helpful assistant with detailed instructions..." * 20
        history = "User: ...\nAssistant: ...\n" * 40
        retrieved = "Chunk content about the topic..." * 100
        reserved_output = 4096

        enc = tiktoken.get_encoding("cl100k_base")
        sys_tokens = len(enc.encode(system_prompt))
        hist_tokens = len(enc.encode(history))
        retr_tokens = len(enc.encode(retrieved))

        used = sys_tokens + hist_tokens + retr_tokens + reserved_output
        remaining = MAX_CONTEXT - used

        print(f"  System prompt:     {sys_tokens:6d} tokens")
        print(f"  History:           {hist_tokens:6d} tokens")
        print(f"  Retrieved chunks:  {retr_tokens:6d} tokens")
        print(f"  Reserved output:   {reserved_output:6d} tokens")
        print(f"  Total used:        {used:6d} tokens")
        print(f"  Remaining:         {remaining:6d} tokens")
        if remaining < 0:
            print("  WARNING: Over context budget! Trim system prompt or history.")
        else:
            print("  OK: Within context budget.")

        # --- 5. Chat message overhead ---
        print("\n=== Chat Template Overhead ===")
        # Even an empty message has token overhead from role markers
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
        ]
        # Rough approximation of chat template overhead
        raw_text = "".join(m["content"] for m in messages)
        raw_tokens = len(enc.encode(raw_text))
        print(f"  Raw content tokens: {raw_tokens}")
        print(f"  (Role markers add a small constant per message — budget for it)")


if __name__ == "__main__":
    main()