"""Token counting utility.

Uses tiktoken (OpenAI's BPE tokeniser) when available.
Falls back to a character-based approximation when tiktoken is absent:
  - CJK characters: 1 char ≈ 1 token
  - Latin/other:    4 chars ≈ 1 token
"""

from __future__ import annotations

import logging

log = logging.getLogger(__name__)

# Module-level caches
_encoder_cache: dict[str, object] = {}
_tiktoken_available: bool | None = None  # None = unchecked


def _check_tiktoken() -> bool:
    global _tiktoken_available
    if _tiktoken_available is None:
        try:
            import tiktoken  # noqa: F401

            _tiktoken_available = True
        except ImportError:
            _tiktoken_available = False
            log.debug("tiktoken not installed — using character approximation")
    return _tiktoken_available


def _get_encoder(encoding_name: str):
    """Return a cached tiktoken encoder, or None on error."""
    if encoding_name not in _encoder_cache:
        import tiktoken

        try:
            _encoder_cache[encoding_name] = tiktoken.get_encoding(encoding_name)
        except Exception as exc:
            log.warning("tiktoken: cannot load encoding %r: %s", encoding_name, exc)
            _encoder_cache[encoding_name] = None
    return _encoder_cache[encoding_name]


def count_tokens(text: str, encoding_name: str = "cl100k_base") -> int:
    """Return approximate token count for *text*.

    Args:
        text:          The string to tokenise.
        encoding_name: tiktoken encoding name (default "cl100k_base", used by
                       GPT-4 and Deepseek). Ignored when falling back.

    Returns:
        Estimated token count (≥ 1 for non-empty strings, 0 for empty).
    """
    if not text:
        return 0

    if _check_tiktoken():
        enc = _get_encoder(encoding_name)
        if enc is not None:
            try:
                return len(enc.encode(text))
            except Exception as exc:
                log.debug("tiktoken encode error: %s — falling back", exc)

    # Fallback: character-based approximation
    return _approx_tokens(text)


def _approx_tokens(text: str) -> int:
    """Rough token estimate without tiktoken.

    CJK characters each map roughly to 1 token; Latin characters ~4 per token.
    """
    cjk = sum(
        1
        for c in text
        if (
            "\u4e00" <= c <= "\u9fff"  # CJK Unified Ideographs
            or "\u3000" <= c <= "\u303f"  # CJK Symbols and Punctuation
            or "\uff00" <= c <= "\uffef"  # Halfwidth and Fullwidth Forms
            or "\u3400" <= c <= "\u4dbf"  # CJK Extension A
        )
    )
    latin = len(text) - cjk
    # Ceiling division for Latin chars; zero contribution when no Latin chars.
    latin_tokens = (latin + 3) // 4 if latin > 0 else 0
    return cjk + latin_tokens
