import time
from functools import wraps

class TokenTracker:
    def __init__(self):
        self.stats = {
            'total_requests': 0,
            'total_input_tokens': 0,
            'total_output_tokens': 0,
            'total_cost': 0,
            'cache_hits': 0
        }
    
    def track(self, model, input_tokens, output_tokens, cached=False):
        """追踪单次调用"""
        pricing = {
            'gpt-4o': {'input': 2.5, 'output': 10.0},      # per 1M tokens
            'gpt-4o-mini': {'input': 0.15, 'output': 0.6},
            'claude-3-sonnet': {'input': 3.0, 'output': 15.0}
        }
        
        model_pricing = pricing.get(model, pricing['gpt-4o'])
        cost = (input_tokens * model_pricing['input'] + 
                output_tokens * model_pricing['output']) / 1_000_000
        
        self.stats['total_requests'] += 1
        self.stats['total_input_tokens'] += input_tokens
        self.stats['total_output_tokens'] += output_tokens
        self.stats['total_cost'] += cost
        
        if cached:
            self.stats['cache_hits'] += 1
        
        return {
            'model': model,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'cost_usd': cost,
            'cached': cached
        }

# 使用示例
tracker = TokenTracker()

# 包装你的API调用
def call_llm_with_tracking(prompt, model='gpt-4o-mini'):
    input_tokens = estimate_tokens(prompt)
    
    # 检查缓存
    cached_response = cache.get(hash(prompt))
    if cached_response:
        tracker.track(model, input_tokens, len(cached_response), cached=True)
        return cached_response
    
    # 实际调用
    response = openai.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )
    
    output_tokens = response.usage.completion_tokens
    tracker.track(model, input_tokens, output_tokens, cached=False)
    
    return response.choices[0].message.content

class RequestCostMonitor:
    """追踪单个请求的成本"""
    
    def __init__(self, request_id):
        self.request_id = request_id
        self.llm_calls = []
    
    def add_call(self, model, input_tokens, output_tokens, purpose):
        self.llm_calls.append({
            'timestamp': time.time(),
            'model': model,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'purpose': purpose  # 标记调用用途，如"intent_classification"
        })
    
    def summary(self):
        total_input = sum(c['input_tokens'] for c in self.llm_calls)
        total_output = sum(c['output_tokens'] for c in self.llm_calls)
        
        # 按用途分组统计
        by_purpose = {}
        for call in self.llm_calls:
            p = call['purpose']
            if p not in by_purpose:
                by_purpose[p] = {'calls': 0, 'tokens': 0}
            by_purpose[p]['calls'] += 1
            by_purpose[p]['tokens'] += call['input_tokens'] + call['output_tokens']
        
        return {
            'request_id': self.request_id,
            'llm_calls_count': len(self.llm_calls),
            'total_input_tokens': total_input,
            'total_output_tokens': total_output,
            'breakdown_by_purpose': by_purpose
        }

# 集成到Web框架
from contextvars import ContextVar

current_request_cost = ContextVar('request_cost', default=None)

@app.middleware("http")
async def cost_tracking_middleware(request, call_next):
    request_id = str(uuid.uuid4())
    cost_monitor = RequestCostMonitor(request_id)
    current_request_cost.set(cost_monitor)
    
    response = await call_next(request)
    
    # 记录成本摘要
    summary = cost_monitor.summary()
    logger.info(f"Request {request_id} cost summary: {summary}")
    
    # 写入成本追踪系统
    cost_exporter.export(summary)
    
    return response

class ConversationCostTracker:
    """追踪对话会话的累计成本"""
    
    def __init__(self, conversation_id):
        self.conversation_id = conversation_id
        self.messages = []
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.model_usage = defaultdict(lambda: {'input': 0, 'output': 0})
    
    def add_turn(self, user_message, assistant_response, model):
        """记录一轮对话"""
        input_tokens = estimate_tokens(user_message)
        output_tokens = estimate_tokens(assistant_response)
        
        self.messages.append({
            'user': user_message,
            'assistant': assistant_response,
            'model': model,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'timestamp': datetime.now()
        })
        
        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens
        self.model_usage[model]['input'] += input_tokens
        self.model_usage[model]['output'] += output_tokens
    
    def get_session_summary(self):
        """获取会话成本摘要"""
        total_cost = self._calculate_cost()
        avg_cost_per_turn = total_cost / max(len(self.messages), 1)
        
        return {
            'conversation_id': self.conversation_id,
            'turns': len(self.messages),
            'total_input_tokens': self.total_input_tokens,
            'total_output_tokens': self.total_output_tokens,
            'total_cost_usd': total_cost,
            'avg_cost_per_turn': avg_cost_per_turn,
            'model_breakdown': dict(self.model_usage),
            'duration_minutes': self._get_duration()
        }
    
    def _calculate_cost(self):
        # 根据实际使用的模型和Token数计算成本
        pass

class BusinessCostAnalyzer:
    """业务层成本分析"""
    
    def analyze_customer_cohort(self, start_date, end_date):
        """分析客户群组的AI成本"""
        query = """
        SELECT 
            customer_id,
            COUNT(DISTINCT conversation_id) as conversation_count,
            SUM(total_cost_usd) as total_cost,
            SUM(total_cost_usd) / COUNT(DISTINCT conversation_id) as cost_per_conversation,
            customer_tier
        FROM conversation_costs
        WHERE date BETWEEN %s AND %s
        GROUP BY customer_id, customer_tier
        """
        
        results = db.execute(query, (start_date, end_date))
        
        # 分析不同客户层级的成本
        tier_analysis = defaultdict(list)
        for row in results:
            tier_analysis[row['customer_tier']].append(row['cost_per_conversation'])
        
        return {
            tier: {
                'avg_cost_per_conversation': sum(costs) / len(costs),
                'total_customers': len(costs),
                'p95_cost': np.percentile(costs, 95)
            }
            for tier, costs in tier_analysis.items()
        }
    
    def calculate_feature_roi(self, feature_name):
        """计算某个AI功能的ROI"""
        # 获取该功能的AI成本
        ai_cost = self.get_feature_ai_cost(feature_name)
        
        # 获取该功能带来的收益（根据业务指标）
        revenue_impact = self.get_feature_revenue_impact(feature_name)
        
        roi = (revenue_impact - ai_cost) / ai_cost if ai_cost > 0 else 0
        
        return {
            'feature': feature_name,
            'ai_cost': ai_cost,
            'revenue_impact': revenue_impact,
            'roi': roi,
            'roi_percentage': f"{roi * 100:.1f}%"
        }

# ❌ 低效：每次构建完整Prompt
prompt = f"""
你是一个智能客服助手。你的任务是帮助用户解决问题。
用户的问题是：{user_question}
请基于以下知识库内容回答：
{large_knowledge_base_content}
请记住要友好、专业、简洁。
"""

# ✅ 高效：使用系统Prompt模板 + 精简用户内容
SYSTEM_PROMPT = """你是智能客服助手。基于知识库回答用户问题，要求：
1. 友好专业
2. 回答简洁
3. 不确定时诚实告知"""

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": f"问题：{user_question}\n参考资料：{relevant_chunk}"}
]

class ContextSelector:
    """智能选择最相关的上下文"""
    
    def __init__(self, max_context_tokens=2000):
        self.max_context_tokens = max_context_tokens
    
    def select_context(self, query, candidate_docs):
        """
        选择最相关的文档，控制在Token预算内
        """
        # 按相关性排序
        sorted_docs = sorted(candidate_docs, 
                            key=lambda x: x.relevance_score, 
                            reverse=True)
        
        selected = []
        total_tokens = estimate_tokens(query)  # 预留query空间
        
        for doc in sorted_docs:
            doc_tokens = estimate_tokens(doc.content)
            
            if total_tokens + doc_tokens <= self.max_context_tokens:
                selected.append(doc)
                total_tokens += doc_tokens
            else:
                # 尝试截断最后一份文档
                remaining = self.max_context_tokens - total_tokens
                if remaining > 100:  # 至少保留100 tokens
                    truncated = self._truncate_to_tokens(doc.content, remaining)
                    selected.append(Document(content=truncated, source=doc.source))
                break
        
        return selected
    
    def _truncate_to_tokens(self, text, max_tokens):
        """按Token截断文本"""
        # 粗略估算：1 token ≈ 0.75个英文单词或0.4个中文字符
        words = text.split()
        estimated_tokens = len(words) / 0.75
        
        if estimated_tokens <= max_tokens:
            return text
        
        # 按比例截断
        ratio = max_tokens / estimated_tokens
        keep_words = int(len(words) * ratio)
        return ' '.join(words[:keep_words]) + "..."

class ConversationCompressor:
    """压缩对话历史，保留关键信息"""
    
    def compress_history(self, messages, max_messages=10):
        """
        压缩策略：
        1. 保留最近的N轮完整对话
        2. 更早的对话做摘要
        3. 系统指令始终保留
        """
        if len(messages) <= max_messages:
            return messages
        
        # 分离系统消息
        system_msgs = [m for m in messages if m['role'] == 'system']
        conversation = [m for m in messages if m['role'] != 'system']
        
        # 最近的完整保留
        recent = conversation[-max_messages:]
        
        # 更早的做摘要
        older = conversation[:-max_messages]
        if older:
            summary = self._summarize_messages(older)
            return system_msgs + [{"role": "system", "content": f"历史摘要：{summary}"}] + recent
        
        return system_msgs + recent
    
    def _summarize_messages(self, messages):
        """使用轻量模型或规则对历史做摘要"""
        # 方案1：用轻量模型
        summary_prompt = "请摘要以下对话的关键信息（50字以内）：\n"
        for m in messages:
            summary_prompt += f"{m['role']}: {m['content']}\n"
        
        # 用gpt-4o-mini做摘要，成本低
        response = call_light_model(summary_prompt)
        return response

# 使用max_tokens限制输出长度
response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    max_tokens=500,  # 明确限制输出长度
    temperature=0.7
)

# 在Prompt中明确要求简洁
prompt = """分析以下代码的问题，回答限制在3句话以内：
{code}
"""

class ModelRouter:
    """基于任务复杂度选择合适模型"""
    
    PRICING = {
        'gpt-4o': {'input': 2.5, 'output': 10.0},
        'gpt-4o-mini': {'input': 0.15, 'output': 0.6},
        'claude-3-haiku': {'input': 0.25, 'output': 1.25}
    }
    
    def __init__(self):
        self.classifier = IntentClassifier()  # 轻量分类器
    
    def route(self, query, context=None):
        """选择最适合的模型"""
        
        # 简单意图分类
        complexity = self.classifier.classify(query)
        
        # 根据复杂度选择模型
        if complexity == 'simple':
            # 简单任务：问候、常见问题、简单分类
            return 'gpt-4o-mini'
        elif complexity == 'medium':
            # 中等任务：需要一定推理但不太复杂
            return 'gpt-4o'
        else:
            # 复杂任务：代码生成、复杂推理、创意写作
            return 'gpt-4o'
    
    def estimate_cost(self, query, model):
        """预估成本"""
        input_tokens = estimate_tokens(query)
        # 假设输出是输入的1.5倍
        estimated_output = int(input_tokens * 1.5)
        
        pricing = self.PRICING[model]
        cost = (input_tokens * pricing['input'] + 
                estimated_output * pricing['output']) / 1_000_000
        
        return {
            'model': model,
            'estimated_input_tokens': input_tokens,
            'estimated_output_tokens': estimated_output,
            'estimated_cost_usd': cost
        }

class BatchProcessor:
    """批量处理请求以降低成本"""
    
    def __init__(self, batch_size=100):
        self.batch_size = batch_size
        self.queue = []
    
    async def add_request(self, request):
        self.queue.append(request)
        
        if len(self.queue) >= self.batch_size:
            await self._flush_batch()
    
    async def _flush_batch(self):
        if not self.queue:
            return
        
        batch = self.queue[:self.batch_size]
        self.queue = self.queue[self.batch_size:]
        
        # 合并相似请求，减少重复计算
        grouped = self._group_similar_requests(batch)
        
        for group in grouped.values():
            # 批量调用API（如果API支持）
            # 或并行调用
            await asyncio.gather(*[
                self._process_single(req) for req in group
            ])
    
    def _group_similar_requests(self, requests):
        """将相似请求分组，可用于缓存优化"""
        groups = defaultdict(list)
        for req in requests:
            # 使用query的hash作为分组键
            key = self._get_similarity_key(req.query)
            groups[key].append(req)
        return groups

import hashlib
from functools import lru_cache

class ExactCache:
    """精确匹配缓存"""
    
    def __init__(self, redis_client, ttl=3600):
        self.redis = redis_client
        self.ttl = ttl
    
    def _get_key(self, prompt, model, params):
        """生成缓存键"""
        key_data = f"{prompt}:{model}:{sorted(params.items())}"
        return f"llm:cache:{hashlib.md5(key_data.encode()).hexdigest()}"
    
    def get(self, prompt, model='gpt-4o-mini', **params):
        key = self._get_key(prompt, model, params)
        cached = self.redis.get(key)
        if cached:
            return json.loads(cached)
        return None
    
    def set(self, prompt, response, model='gpt-4o-mini', **params):
        key = self._get_key(prompt, model, params)
        self.redis.setex(
            key, 
            self.ttl, 
            json.dumps({'response': response, 'cached_at': time.time()})
        )

from sentence_transformers import SentenceTransformer
import numpy as np

class SemanticCache:
    """语义相似度缓存"""
    
    def __init__(self, vector_db, similarity_threshold=0.95):
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.vector_db = vector_db
        self.threshold = similarity_threshold
    
    def get(self, query):
        """查找语义相似的缓存结果"""
        # 编码查询
        query_embedding = self.encoder.encode(query)
        
        # 向量检索
        results = self.vector_db.similarity_search(
            embedding=query_embedding,
            k=1,
            filter={'type': 'llm_response'}
        )
        
        if results and results[0]['score'] >= self.threshold:
            return {
                'hit': True,
                'similarity': results[0]['score'],
                'cached_query': results[0]['metadata']['query'],
                'response': results[0]['metadata']['response']
            }
        
        return {'hit': False}
    
    def set(self, query, response):
        """缓存新的查询-响应对"""
        embedding = self.encoder.encode(query)
        
        self.vector_db.add_document({
            'embedding': embedding,
            'metadata': {
                'type': 'llm_response',
                'query': query,
                'response': response,
                'timestamp': time.time()
            }
        })

class CacheabilityAnalyzer:
    """分析查询是否值得缓存"""
    
    def __init__(self):
        # 高频查询模式
        self.high_volume_patterns = [
            r'.*什么是.*',
            r'.*怎么.*',
            r'.*为什么.*',
            r'.*价格.*',
            r'.*(你好|您好|hi|hello).*'
        ]
        
        # 不应该缓存的模式（个性化、时效性）
        self.no_cache_patterns = [
            r'.*我的.*订单.*',
            r'.*今天.*',
            r'.*最新.*',
            r'.*现在.*'
        ]
    
    def should_cache(self, query):
        """判断是否应该缓存该查询"""
        
        # 检查是否应该排除
        for pattern in self.no_cache_patterns:
            if re.match(pattern, query, re.IGNORECASE):
                return False, 'personalized_or_time_sensitive'
        
        # 检查是否是高价值缓存目标
        for pattern in self.high_volume_patterns:
            if re.match(pattern, query, re.IGNORECASE):
                return True, 'high_volume_pattern'
        
        # 默认：中等长度、通用性查询适合缓存
        if len(query) < 200 and not any(word in query for word in ['我', '我的']):
            return True, 'general_query'
        
        return False, 'low_cache_value'

class SmartRouter:
    """智能路由：根据查询特征选择最优路径"""
    
    def __init__(self):
        self.exact_cache = ExactCache(redis_client)
        self.semantic_cache = SemanticCache(vector_db)
        self.model_router = ModelRouter()
        self.cost_tracker = CostTracker()
    
    async def route(self, query, context=None, min_quality='medium'):
        """
        智能路由流程
        
        路径优先级：
        1. 精确缓存（成本≈0，延迟<1ms）
        2. 语义缓存（成本≈0，延迟<100ms）
        3. 轻量模型（成本低，质量中等）
        4. 强模型（成本高，质量高）
        """
        start_time = time.time()
        
        # 1. 检查精确缓存
        cached = self.exact_cache.get(query)
        if cached:
            self.cost_tracker.record_cache_hit('exact')
            return {
                'response': cached['response'],
                'source': 'exact_cache',
                'cost': 0,
                'latency_ms': (time.time() - start_time) * 1000
            }
        
        # 2. 检查语义缓存
        semantic_hit = self.semantic_cache.get(query)
        if semantic_hit['hit']:
            self.cost_tracker.record_cache_hit('semantic')
            return {
                'response': semantic_hit['response'],
                'source': 'semantic_cache',
                'similarity': semantic_hit['similarity'],
                'cost': 0,
                'latency_ms': (time.time() - start_time) * 1000
            }
        
        # 3. 选择模型
        model = self.model_router.route(query, context)
        
        # 如果质量要求低，强制用轻量模型
        if min_quality == 'low':
            model = 'gpt-4o-mini'
        
        # 4. 调用模型
        response = await self._call_model(query, model, context)
        
        # 5. 更新缓存
        if self._should_cache_response(query, response):
            self.exact_cache.set(query, response, model)
            self.semantic_cache.set(query, response)
        
        cost = self.cost_tracker.calculate_cost(query, response, model)
        
        return {
            'response': response,
            'source': f'model:{model}',
            'cost': cost,
            'latency_ms': (time.time() - start_time) * 1000
        }
    
    def _should_cache_response(self, query, response):
        """判断响应是否值得缓存"""
        # 错误响应不缓存
        if 'error' in response.lower():
            return False
        
        # 太短的响应可能价值不大
        if len(response) < 20:
            return False
        
        return True

class ValueBasedRouter:
    """基于业务价值的路由决策"""
    
    def __init__(self):
        self.pricing = {
            'gpt-4o-mini': 0.6,      # per 1K output tokens ($)
            'gpt-4o': 10.0,
            'claude-3-opus': 75.0
        }
    
    def decide(self, query_context):
        """
        基于价值判断选择模型
        
        决策因素：
        1. 用户价值（VIP客户？）
        2. 任务关键度（是否影响交易？）
        3. 错误成本（出错的后果？）
        4. 时间敏感度（是否需要立即回答？）
        """
        
        user_value = query_context.get('user_ltv', 0)  # 用户生命周期价值
        task_criticality = query_context.get('criticality', 'low')
        error_cost = query_context.get('error_cost', 0)
        
        # 计算"质量溢价"的价值
        quality_premium_value = self._calculate_quality_value(
            user_value, task_criticality, error_cost
        )
        
        # 模型成本差异
        mini_cost = self._estimate_cost('gpt-4o-mini', query_context)
        full_cost = self._estimate_cost('gpt-4o', query_context)
        
        cost_difference = full_cost - mini_cost
        
        # 决策：如果质量价值 > 成本差异，用强模型
        if quality_premium_value > cost_difference * 10:  # 10x安全边际
            return 'gpt-4o', 'high_quality_warranted'
        else:
            return 'gpt-4o-mini', 'cost_optimized'
    
    def _calculate_quality_value(self, user_value, criticality, error_cost):
        """计算更高质量输出的业务价值"""
        
        # 基础价值
        base_value = user_value * 0.01  # 假设好的体验带来1% LTV提升
        
        # 关键度调整
        criticality_multiplier = {
            'critical': 5.0,
            'high': 2.0,
            'medium': 1.0,
            'low': 0.5
        }.get(criticality, 1.0)
        
        # 错误成本
        error_prevention_value = error_cost * 0.2  # 强模型减少20%错误
        
        return (base_value * criticality_multiplier) + error_prevention_value

class CostQualityExperiment:
    """成本-质量权衡的A/B测试框架"""
    
    def run_experiment(self, traffic_split=0.1):
        """
        对比两种策略的业务效果
        """
        
        # 对照组：成本优先
        def control_group(query):
            return self.router.route(query, strategy='cost_first')
        
        # 实验组：质量优先
        def treatment_group(query):
            return self.router.route(query, strategy='quality_first')
        
        # 分流
        if random.random() < traffic_split:
            # 实验组
            result = treatment_group(query)
            result['experiment_group'] = 'treatment'
        else:
            # 对照组
            result = control_group(query)
            result['experiment_group'] = 'control'
        
        return result
    
    def analyze_results(self, experiment_data):
        """分析实验结果"""
        
        control = [d for d in experiment_data if d['group'] == 'control']
        treatment = [d for d in experiment_data if d['group'] == 'treatment']
        
        metrics = {
            'control': {
                'avg_cost': sum(d['cost'] for d in control) / len(control),
                'satisfaction': sum(d['rating'] for d in control) / len(control),
                'conversion': sum(d['converted'] for d in control) / len(control)
            },
            'treatment': {
                'avg_cost': sum(d['cost'] for d in treatment) / len(treatment),
                'satisfaction': sum(d['rating'] for d in treatment) / len(treatment),
                'conversion': sum(d['converted'] for d in treatment) / len(treatment)
            }
        }
        
        # 计算ROI
        cost_increase = metrics['treatment']['avg_cost'] - metrics['control']['avg_cost']
        conversion_lift = metrics['treatment']['conversion'] - metrics['control']['conversion']
        
        roi = conversion_lift / cost_increase if cost_increase > 0 else 0
        
        return {
            'metrics': metrics,
            'cost_increase_per_request': cost_increase,
            'conversion_lift': conversion_lift,
            'roi': roi
        }