"""
Tiktoken Token计数插件
支持OpenAI和其他使用tiktoken的模型
"""

import json
from typing import Any, Dict, List, Optional

from src.core.logger import logger

from .base import TokenCounterPlugin, TokenUsage

# 尝试导入tiktoken
try:
    import tiktoken

    TIKTOKEN_AVAILABLE = True
except ImportError:
    TIKTOKEN_AVAILABLE = False
    tiktoken = None


class TiktokenCounterPlugin(TokenCounterPlugin):
    """
    使用tiktoken库计算Token数量
    支持OpenAI模型和其他兼容模型
    """

    # 模型编码映射
    MODEL_ENCODINGS = {
        # GPT-4 系列
        "gpt-4": "cl100k_base",
        "gpt-4-32k": "cl100k_base",
        "gpt-4-turbo": "cl100k_base",
        "gpt-4-turbo-preview": "cl100k_base",
        "gpt-4o": "o200k_base",
        "gpt-4o-mini": "o200k_base",
        # GPT-3.5 系列
        "gpt-3.5-turbo": "cl100k_base",
        "gpt-3.5-turbo-16k": "cl100k_base",
        # 旧模型
        "text-davinci-003": "p50k_base",
        "text-davinci-002": "p50k_base",
        "code-davinci-002": "p50k_base",
        # Embeddings
        "text-embedding-ada-002": "cl100k_base",
        "text-embedding-3-small": "cl100k_base",
        "text-embedding-3-large": "cl100k_base",
    }

    # 每个消息的额外Token数
    MESSAGE_OVERHEAD = {
        "gpt-3.5-turbo": 4,  # 每条消息
        "gpt-4": 3,
        "gpt-4-turbo": 3,
        "gpt-4o": 3,
        "gpt-4o-mini": 3,
    }

    def __init__(self, name: str = "tiktoken", config: Dict[str, Any] = None):
        super().__init__(name, config)

        if not TIKTOKEN_AVAILABLE:
            self.enabled = False
            logger.warning("tiktoken not installed, plugin disabled")
            return

        # 缓存编码器
        self._encoders = {}

        # 价格表（每1M tokens的价格 USD）
        default_pricing = {
            "gpt-4o": {"input": 2.5, "output": 10},
            "gpt-4o-mini": {"input": 0.15, "output": 0.6},
            "gpt-4-turbo": {"input": 10, "output": 30},
            "gpt-4": {"input": 30, "output": 60},
            "gpt-3.5-turbo": {"input": 0.5, "output": 1.5},
            "o1-preview": {"input": 15, "output": 60, "reasoning": 60},
            "o1-mini": {"input": 3, "output": 12, "reasoning": 12},
        }
        self.config["pricing"] = (
            config.get("pricing", default_pricing) if config else default_pricing
        )

    def _get_encoder(self, model: str) -> Any:
        """获取模型的编码器"""
        if model in self._encoders:
            return self._encoders[model]

        # 获取编码名称
        encoding_name = None

        # 完全匹配
        if model in self.MODEL_ENCODINGS:
            encoding_name = self.MODEL_ENCODINGS[model]
        else:
            # 前缀匹配
            for model_prefix, enc_name in self.MODEL_ENCODINGS.items():
                if model.startswith(model_prefix):
                    encoding_name = enc_name
                    break

        # 如果找不到，尝试使用模型名称
        if not encoding_name:
            try:
                encoder = tiktoken.encoding_for_model(model)
                self._encoders[model] = encoder
                return encoder
            except:
                # 默认使用cl100k_base
                encoding_name = "cl100k_base"

        # 创建编码器
        encoder = tiktoken.get_encoding(encoding_name)
        self._encoders[model] = encoder
        return encoder

    def supports_model(self, model: str) -> bool:
        """检查是否支持指定模型"""
        # 支持所有OpenAI模型和一些兼容模型
        openai_models = ["gpt-4", "gpt-3.5", "text-davinci", "text-embedding", "code-davinci", "o1"]
        return any(model.startswith(prefix) for prefix in openai_models)

    async def count_tokens(self, text: str, model: Optional[str] = None) -> int:
        """计算文本的Token数量"""
        if not self.enabled:
            return 0

        model = model or self.default_model or "gpt-3.5-turbo"
        encoder = self._get_encoder(model)

        try:
            tokens = encoder.encode(text)
            return len(tokens)
        except Exception as e:
            logger.warning(f"Error counting tokens: {e}")
            # 简单估算: 平均每个字符0.75个token
            return int(len(text) * 0.75)

    async def count_messages(
        self, messages: List[Dict[str, Any]], model: Optional[str] = None
    ) -> int:
        """计算消息列表的Token数量"""
        if not self.enabled:
            return 0

        model = model or self.default_model or "gpt-3.5-turbo"
        encoder = self._get_encoder(model)

        # 获取每条消息的额外token数
        msg_overhead = self.MESSAGE_OVERHEAD.get(model, 3)

        total_tokens = 0

        for message in messages:
            # 每条消息的基本token
            total_tokens += msg_overhead

            # 角色token
            role = message.get("role", "")
            if role:
                total_tokens += len(encoder.encode(role))

            # 内容token
            content = message.get("content")
            if content:
                if isinstance(content, str):
                    total_tokens += len(encoder.encode(content))
                elif isinstance(content, list):
                    # 处理多模态内容
                    for item in content:
                        if item.get("type") == "text":
                            text = item.get("text", "")
                            total_tokens += len(encoder.encode(text))
                        elif item.get("type") == "image_url":
                            # 图像的token计算更复杂，这里简化处理
                            # 低分辨率: 85 tokens, 高分辨率: 170 tokens
                            detail = item.get("image_url", {}).get("detail", "auto")
                            total_tokens += 170 if detail == "high" else 85

            # 名称token
            name = message.get("name")
            if name:
                total_tokens += len(encoder.encode(name)) - 1  # name会减去1个token

            # 工具调用
            tool_calls = message.get("tool_calls")
            if tool_calls:
                for tool_call in tool_calls:
                    # 工具ID
                    if "id" in tool_call:
                        total_tokens += len(encoder.encode(tool_call["id"]))

                    # 函数信息
                    function = tool_call.get("function", {})
                    if "name" in function:
                        total_tokens += len(encoder.encode(function["name"]))
                    if "arguments" in function:
                        total_tokens += len(encoder.encode(function["arguments"]))

        # 添加固定的结束标记
        total_tokens += 3

        return total_tokens

    async def get_model_info(self, model: str) -> Dict[str, Any]:
        """获取模型信息"""
        info = {"model": model, "supported": self.supports_model(model)}

        if self.supports_model(model):
            # 获取编码信息
            encoder = self._get_encoder(model)
            encoding_name = None

            # 找到编码名称
            for m, enc in self.MODEL_ENCODINGS.items():
                if model.startswith(m):
                    encoding_name = enc
                    break

            info.update(
                {
                    "encoding": encoding_name or "unknown",
                    "vocab_size": encoder.n_vocab if hasattr(encoder, "n_vocab") else None,
                    "max_tokens": self._get_max_tokens(model),
                    "message_overhead": self.MESSAGE_OVERHEAD.get(model, 3),
                }
            )

            # 添加价格信息
            pricing = self.config.get("pricing", {})
            if model in pricing:
                info["pricing"] = pricing[model]

        return info

    def _get_max_tokens(self, model: str) -> int:
        """获取模型的最大token数"""
        max_tokens_map = {
            "gpt-4": 8192,
            "gpt-4-32k": 32768,
            "gpt-4-turbo": 128000,
            "gpt-4o": 128000,
            "gpt-4o-mini": 128000,
            "gpt-3.5-turbo": 4096,
            "gpt-3.5-turbo-16k": 16384,
            "o1-preview": 128000,
            "o1-mini": 128000,
        }

        # 完全匹配
        if model in max_tokens_map:
            return max_tokens_map[model]

        # 前缀匹配
        for model_prefix, max_tokens in max_tokens_map.items():
            if model.startswith(model_prefix):
                return max_tokens

        # 默认值
        return 4096

    async def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = await super().get_stats()
        stats.update(
            {"encoders_cached": len(self._encoders), "tiktoken_available": TIKTOKEN_AVAILABLE}
        )
        return stats
