from __future__ import annotations

from typing import Callable, Optional


def _fallback_count_tokens(text: str) -> int:
    if not text:
        return 0
    return len(text) // 4


def get_tokenizer(model: Optional[str] = None) -> Callable[[str], int]:
    try:
        import tiktoken  # type: ignore

        try:
            if model:
                encoding = tiktoken.encoding_for_model(model)
            else:
                encoding = tiktoken.get_encoding("cl100k_base")
        except Exception:
            encoding = tiktoken.get_encoding("cl100k_base")

        def _count_with_encoding(text: str) -> int:
            return len(encoding.encode(text or ""))

        return _count_with_encoding
    except Exception:
        return _fallback_count_tokens


def count_tokens(text: str, model: Optional[str] = None) -> int:
    tokenizer = get_tokenizer(model=model)
    return tokenizer(text)
