import tiktoken # OpenAI的Tokenizer库，也常用于其他模型估算
from dataclasses import dataclass
from typing import Dict

@dataclass
class ModelPricing:
    """模型定价配置"""
    name: str
    input_price_per_million: float  # 美元/百万Tokens
    output_price_per_million: float # 美元/百万Tokens

class CostCalculator:
    """成本计算与监控器"""
    
    # 预定义一些模型价格（示例值，需实时更新）
    PRICING_MAP: Dict[str, ModelPricing] = {
        "gpt-4o": ModelPricing("gpt-4o", 5.0, 15.0),
        "gpt-3.5-turbo": ModelPricing("gpt-3.5-turbo", 0.5, 1.5),
        "claude-3-haiku": ModelPricing("claude-3-haiku", 0.25, 1.25),
    }
    
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.pricing = self.PRICING_MAP.get(model_name)
        if not self.pricing:
            raise ValueError(f"未知模型定价: {model_name}")
        # 初始化编码器（此处以cl100k_base为例，实际需匹配模型）
        self.encoder = tiktoken.get_encoding("cl100k_base")
    
    def count_tokens(self, text: str) -> int:
        """计算文本的Token数量"""
        return len(self.encoder.encode(text))
    
    def calculate_cost(self, prompt: str, completion: str) -> Dict:
        """计算单次调用的成本和Token消耗"""
        input_tokens = self.count_tokens(prompt)
        output_tokens = self.count_tokens(completion)
        
        input_cost = (input_tokens / 1_000_000) * self.pricing.input_price_per_million
        output_cost = (output_tokens / 1_000_000) * self.pricing.output_price_per_million
        total_cost = input_cost + output_cost
        
        return {
            "model": self.model_name,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "input_cost_usd": round(input_cost, 6),
            "output_cost_usd": round(output_cost, 6),
            "total_cost_usd": round(total_cost, 6)
        }
    
    def log_usage(self, prompt: str, completion: str, user_id: str = None):
        """记录使用情况，可用于监控和预算告警"""
        cost_info = self.calculate_cost(prompt, completion)
        # 这里可以集成到日志系统（如ELK）、监控系统（如Prometheus）或数据库
        print(f"[Cost Log] User: {user_id}, Model: {cost_info['model']}, "
              f"Tokens: {cost_info['total_tokens']}, Cost: ${cost_info['total_cost_usd']:.6f}")
        # 模拟：检查预算阈值并触发告警
        if cost_info['total_cost_usd'] > 0.01: # 假设单次调用预算阈值
            self._trigger_budget_alert(user_id, cost_info['total_cost_usd'])
        return cost_info
    
    def _trigger_budget_alert(self, user_id: str, cost: float):
        """触发预算告警（集成到邮件、Slack等）"""
        print(f"⚠️  Budget Alert for user {user_id}: Single call cost ${cost:.6f} exceeded threshold.")

# 使用示例
if __name__ == "__main__":
    calculator = CostCalculator("gpt-3.5-turbo")
    
    sample_prompt = "请总结以下文章的核心观点：" + "AI是未来。" * 100  # 模拟长提示
    sample_completion = "文章的核心观点是AI技术将在未来社会发展中扮演至关重要的角色。" * 10
    
    result = calculator.log_usage(sample_prompt, sample_completion, user_id="test_user_001")
    print("\n详细成本分析:")
    for key, value in result.items():
        print(f"  {key}: {value}")

# 示例：使用小模型进行文档摘要，以减少主模型的输入长度
import openai # 假设已安装和配置

def summarize_for_context(long_document: str, max_summary_tokens: int = 500) -> str:
    """
    使用廉价模型生成摘要，作为主模型的输入上下文。
    这是一种“两阶段”处理策略。
    """
    client = openai.OpenAI()
    
    summary_prompt = f"""
    请为以下文档生成一个简洁的摘要，摘要长度不超过{max_summary_tokens}个tokens。
    摘要需保留核心事实、数据和结论。
    
    文档：
    {long_document[:10000]}  # 防止过长，可分段处理
    """
    
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",  # 使用廉价模型
            messages=[{"role": "user", "content": summary_prompt}],
            max_tokens=max_summary_tokens,
            temperature=0.1
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        # 降级策略：如果摘要失败，返回文档开头部分
        print(f"摘要生成失败: {e}")
        return long_document[:2000]  # 截断

# 在主流程中使用
long_doc = "..." # 你的长文档
query = "根据文档，AI成本优化的关键是什么？"

# 传统方式：将整个长文档作为上下文
# expensive_prompt = f"文档：{long_doc}\n\n问题：{query}"

# 优化方式：先摘要
summary = summarize_for_context(long_doc)
optimized_prompt = f"基于以下摘要回答问题：\n{summary}\n\n问题：{query}"

# 现在可以将optimized_prompt发送给更强大的主模型（如GPT-4），输入Token大大减少

# 简化版语义缓存实现思路
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import Tuple, Optional
import hashlib
import json

class SemanticCache:
    def __init__(self, similarity_threshold: float = 0.9):
        # 使用轻量级句子嵌入模型
        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
        self.threshold = similarity_threshold
        self.cache = {}  # 实际生产中使用Redis或Memcached
        # key: 嵌入向量的哈希或ID, value: {"response": ..., "metadata": ...}
    
    def _get_embedding(self, text: str) -> np.ndarray:
        """生成文本的嵌入向量"""
        return self.embedder.encode(text, normalize_embeddings=True)
    
    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """计算余弦相似度"""
        return np.dot(vec1, vec2)  # 向量已归一化
    
    def _generate_cache_key(self, embedding: np.ndarray) -> str:
        """将嵌入向量转换为缓存键（简化版，生产环境需更健壮）"""
        # 取前几位浮点数生成哈希
        vec_str = ','.join([f"{x:.4f}" for x in embedding[:10]])
        return hashlib.md5(vec_str.encode()).hexdigest()
    
    def get(self, query: str) -> Tuple[bool, Optional[str]]:
        """
        检查缓存。
        返回: (是否命中, 缓存响应或None)
        """
        query_embedding = self._get_embedding(query)
        
        # 在生产环境中，这里应使用向量数据库进行近似最近邻搜索
        # 此处为简化演示，遍历缓存（不适用于大规模缓存）
        for cache_key, cache_data in self.cache.items():
            # 假设cache_data中存储了 embedding
            cached_embedding = cache_data.get("embedding")
            if cached_embedding is not None:
                similarity = self._cosine_similarity(query_embedding, cached_embedding)
                if similarity >= self.threshold:
                    print(f"语义缓存命中！相似度: {similarity:.3f}")
                    return True, cache_data["response"]
        
        return False, None
    
    def set(self, query: str, response: str):
        """将查询和响应存入缓存"""
        query_embedding = self._get_embedding(query)
        cache_key = self._generate_cache_key(query_embedding)
        
        self.cache[cache_key] = {
            "embedding": query_embedding,
            "response": response,
            "query": query  # 用于调试
        }
        print(f"已缓存查询: {query[:50]}...")

# 使用示例
cache = SemanticCache(similarity_threshold=0.92)

user_queries = [
    "解释一下机器学习中的过拟合现象。",
    "什么是机器学习中的过拟合？",
    "过拟合在机器学习中是什么意思？", # 与第一个语义高度相似
    "如何训练一个神经网络？" # 不相似
]

cached_response_for_overfitting = "过拟合是指模型在训练数据上表现很好，但在未见过的测试数据上表现较差的现象，通常因为模型过于复杂，学习了训练数据中的噪声和细节。"

# 模拟第一次查询，未命中，调用LLM并缓存
hit, response = cache.get(user_queries[0])
if not hit:
    print("缓存未命中，调用LLM...")
    # simulated_llm_response = call_llm(user_queries[0])
    simulated_llm_response = cached_response_for_overfitting
    cache.set(user_queries[0], simulated_llm_response)
    response = simulated_llm_response
print(f"Q: {user_queries[0]}\nA: {response}\n")

# 模拟第二次高度相似的查询，应命中缓存
hit, response = cache.get(user_queries[1])
if hit:
    print(f"Q: {user_queries[1]}\nA: {response}\n")
else:
    print("未命中（理论上应命中，可能阈值设置或嵌入模型导致）")

# 示例：基于规则和轻量级分类的路由器
class ModelRouter:
    def __init__(self):
        # 可以集成一个轻量级文本分类模型来判断意图复杂度
        # 此处使用简单规则演示
        self.simple_keywords = ["你好", "谢谢", "再见", "时间", "天气"]
        self.complex_keywords = ["分析", "对比", "总结", "解释原理", "编写代码"]
    
    def route(self, query: str, history: list = None) -> str:
        """
        根据查询内容决定使用哪个模型。
        返回模型标识符。
        """
        query_lower = query.lower()
        
        # 规则1：非常简短的社交或简单查询 -> 最廉价模型
        if len(query) < 10 and any(kw in query_lower for kw in ["你好", "hi", "ok"]):
            return "claude-3-haiku"  # 或 gpt-3.5-turbo
        
        # 规则2：包含明确简单关键词 -> 小模型
        if any(kw in query_lower for kw in self.simple_keywords):
            return "gpt-3.5-turbo"
        
        # 规则3：包含复杂任务关键词或长文本 -> 大模型
        if len(query) > 100 or any(kw in query_lower for kw in self.complex_keywords):
            return "gpt-4o"
        
        # 规则4：默认使用平衡型模型
        return "claude-3-sonnet"  # 或 gpt-4-turbo
    
    # 更高级的实现：使用一个小型ML模型（如fastText）实时预测查询的“复杂度分数”

# 在API网关或主服务逻辑中集成路由
router = ModelRouter()
user_query = "请对比Transformer和CNN在图像处理任务上的优劣，并给出代码示例。"
selected_model = router.route(user_query)
print(f"查询: '{user_query}'")
print(f"路由决策: 使用模型 -> {selected_model}")

# 然后，调用对应的模型服务
# if selected_model == "gpt-4o":
#     response = call_openai_gpt4o(user_query)
# elif selected_model == "gpt-3.5-turbo":
#     response = call_openai_gpt35(user_query)
# ...