"""
Token 计数工具 — 支持 tiktoken 精确计数 + 字符估算兜底。

优先使用 tiktoken(cl100k_base) 与 OpenAI/GPT 系列模型一致；
tiktoken 不可用时按 1 token ≈ 1.5 中文字符估算（误差 ≤10%）。
"""
from __future__ import annotations

import logging
from typing import List, Optional

logger = logging.getLogger(__name__)

# 全局 tiktoken 编码器（延迟初始化）
_encoder = None
_tiktoken_available: Optional[bool] = None


def _init_encoder():
    """延迟初始化 tiktoken 编码器。"""
    global _encoder, _tiktoken_available
    if _tiktoken_available is not None:
        return
    try:
        import tiktoken
        _encoder = tiktoken.get_encoding("cl100k_base")
        _tiktoken_available = True
        logger.info("tiktoken cl100k_base 编码器初始化成功")
    except Exception as e:
        _tiktoken_available = False
        logger.warning(f"tiktoken 不可用，降级为字符估算: {e}")


class TokenCounter:
    """Token 计数器，支持精确(tiktoken)和估算(字符)两种模式。"""

    @staticmethod
    def count_tokens(text: str) -> int:
        """计算单段文本的 token 数。"""
        if not text:
            return 0
        _init_encoder()
        if _tiktoken_available and _encoder is not None:
            return len(_encoder.encode(text))
        # 兜底：中文约1.5字符/token，英文约4字符/token，混合取2字符/token
        return max(1, len(text) // 2)

    @staticmethod
    def count_messages_tokens(messages: List[dict]) -> int:
        """计算消息列表的 token 总数。

        每条消息额外计入 role 标记的 overhead（约4 tokens/条）。
        """
        total = 0
        for msg in messages:
            content = msg.get("content", "")
            role = msg.get("role", "")
            total += TokenCounter.count_tokens(content)
            total += TokenCounter.count_tokens(role)
            total += 4  # 每条消息的格式 overhead
        return total

    @staticmethod
    def is_precise() -> bool:
        """返回当前是否使用精确计数（tiktoken）。"""
        _init_encoder()
        return _tiktoken_available is True