import requests
import time
import statistics
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Tuple
import json


class EmbeddingPerformanceTester:
    def __init__(self, base_url: str = "http://127.0.0.1:9997/v1"):
        """初始化性能测试器"""
        self.base_url = base_url
        self.embeddings_url = f"{base_url}/embeddings"
        
    def test_single_request(self, model: str, input_text: List[str]) -> Tuple[bool, float]:
        """测试单个请求，返回成功状态和耗时"""
        try:
            start_time = time.perf_counter()
            
            # 构建请求体
            # 如果 input_text 是列表，取第一个元素；如果是字符串，直接使用
            input_value = input_text[0] if isinstance(input_text, list) and len(input_text) > 0 else input_text
            
            response = requests.post(
                self.embeddings_url,
                headers={
                    'accept': 'application/json',
                    'Content-Type': 'application/json'
                },
                json={
                    "model": model,
                    "input": input_value
                },
                timeout=60  # 设置超时时间
            )
            
            end_time = time.perf_counter()
            
            # 验证响应格式
            if response.status_code == 200:
                result = response.json()
                if result and 'data' in result and len(result['data']) > 0:
                    return True, end_time - start_time
                else:
                    return False, end_time - start_time
            else:
                return False, end_time - start_time
                
        except Exception as e:
            print(f"请求失败 - 模型 {model}: {str(e)}")
            return False, 0.0
    
    def test_model_sequential(self, model: str, input_text: List[str], 
                             iterations: int = 1000) -> Dict:
        """顺序执行性能测试"""
        print(f"\n开始顺序测试模型: {model}")
        print(f"测试次数: {iterations}")
        
        successes = 0
        failures = 0
        latencies = []
        
        for i in range(iterations):
            if i % 100 == 0 and i > 0:
                print(f"  已完成 {i}/{iterations} 次请求...")
                
            success, latency = self.test_single_request(model, input_text)
            
            if success:
                successes += 1
                latencies.append(latency)
            else:
                failures += 1
        
        return self._calculate_stats(model, successes, failures, latencies)
    
    def test_model_concurrent(self, model: str, input_text: List[str],
                             iterations: int = 1000, max_workers: int = 10) -> Dict:
        """并发执行性能测试"""
        print(f"\n开始并发测试模型: {model}")
        print(f"测试次数: {iterations}, 并发数: {max_workers}")
        
        successes = 0
        failures = 0
        latencies = []
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            future_to_request = {
                executor.submit(self.test_single_request, model, input_text): i
                for i in range(iterations)
            }
            
            # 收集结果
            completed = 0
            for future in as_completed(future_to_request):
                completed += 1
                if completed % 100 == 0:
                    print(f"  已完成 {completed}/{iterations} 次请求...")
                
                try:
                    success, latency = future.result()
                    if success:
                        successes += 1
                        latencies.append(latency)
                    else:
                        failures += 1
                except Exception as e:
                    print(f"请求异常: {str(e)}")
                    failures += 1
        
        return self._calculate_stats(model, successes, failures, latencies)
    
    def _calculate_stats(self, model: str, successes: int, 
                        failures: int, latencies: List[float]) -> Dict:
        """计算性能统计信息"""
        if not latencies:
            return {
                "model": model,
                "total_requests": successes + failures,
                "successful_requests": successes,
                "failed_requests": failures,
                "success_rate": 0.0,
                "error": "无成功请求"
            }
        
        stats = {
            "model": model,
            "total_requests": successes + failures,
            "successful_requests": successes,
            "failed_requests": failures,
            "success_rate": successes / (successes + failures) * 100,
            "total_time": sum(latencies),
            "avg_latency": statistics.mean(latencies),
            "min_latency": min(latencies),
            "max_latency": max(latencies),
            "p50_latency": statistics.median(latencies),
            "p95_latency": sorted(latencies)[int(len(latencies) * 0.95)],
            "p99_latency": sorted(latencies)[int(len(latencies) * 0.99)],
            "requests_per_second": len(latencies) / sum(latencies) if sum(latencies) > 0 else 0
        }
        
        # 添加标准差（如果有多于一个样本）
        if len(latencies) > 1:
            stats["std_dev"] = statistics.stdev(latencies)
        
        return stats
    
    def print_results(self, results: Dict):
        """打印测试结果"""
        print("\n" + "="*60)
        print(f"性能测试结果 - {results['model']}")
        print("="*60)
        
        if "error" in results:
            print(f"错误: {results['error']}")
            return
            
        print(f"总请求数: {results['total_requests']}")
        print(f"成功请求: {results['successful_requests']}")
        print(f"失败请求: {results['failed_requests']}")
        print(f"成功率: {results['success_rate']:.2f}%")
        print(f"总耗时: {results['total_time']:.4f}秒")
        print(f"平均延迟: {results['avg_latency']:.4f}秒")
        print(f"最小延迟: {results['min_latency']:.4f}秒")
        print(f"最大延迟: {results['max_latency']:.4f}秒")
        print(f"P50延迟: {results['p50_latency']:.4f}秒")
        print(f"P95延迟: {results['p95_latency']:.4f}秒")
        print(f"P99延迟: {results['p99_latency']:.4f}秒")
        
        if "std_dev" in results:
            print(f"标准差: {results['std_dev']:.4f}秒")
        
        print(f"QPS: {results['requests_per_second']:.2f} 请求/秒")
        print("="*60)
    
    def save_results(self, results_list: List[Dict], filename: str = "performance_results.json"):
        """保存测试结果到JSON文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(results_list, f, indent=2, ensure_ascii=False)
        print(f"\n结果已保存到: {filename}")


def main():
    """主函数"""
    # 初始化测试器
    tester = EmbeddingPerformanceTester()
    
    # 测试配置
    test_input = ["What is the capital of China?"]
    iterations = 1000
    test_models = ['bge-m3', 'Qwen3-Embedding-0.6B']
    
    print("="*60)
    print("Embedding API 性能测试 (HTTP)")
    print("="*60)
    
    all_results = []
    
    # 测试模式选择
    print("\n选择测试模式:")
    print("1. 顺序测试 (Sequential)")
    print("2. 并发测试 (Concurrent)")
    print("3. 两种模式都测试")
    
    mode = input("请输入选择 (1/2/3, 默认1): ").strip()
    
    for model in test_models:
        print(f"\n{'='*60}")
        print(f"测试模型: {model}")
        print(f"{'='*60}")
        
        if mode in ['2', '3']:
            # 并发测试
            concurrent_results = tester.test_model_concurrent(
                model=model,
                input_text=test_input,
                iterations=iterations,
                max_workers=10  # 可根据需要调整并发数
            )
            tester.print_results(concurrent_results)
            concurrent_results["test_mode"] = "concurrent"
            all_results.append(concurrent_results)
        
        if mode in ['1', '3'] or not mode:
            # 顺序测试
            sequential_results = tester.test_model_sequential(
                model=model,
                input_text=test_input,
                iterations=iterations
            )
            tester.print_results(sequential_results)
            sequential_results["test_mode"] = "sequential"
            all_results.append(sequential_results)
    
    # 保存结果
    tester.save_results(all_results)
    
    # 汇总对比
    print("\n" + "="*60)
    print("性能测试汇总对比")
    print("="*60)
    
    for result in all_results:
        if "error" not in result:
            print(f"\n模型: {result['model']} ({result['test_mode']})")
            print(f"  QPS: {result['requests_per_second']:.2f}")
            print(f"  平均延迟: {result['avg_latency']:.4f}秒")
            print(f"  成功率: {result['success_rate']:.2f}%")


if __name__ == "__main__":
    # 添加一个简单的健康检查
    try:
        tester = EmbeddingPerformanceTester()
        # 快速测试连接
        test_result = tester.test_single_request('bge-m3', ["test"])
        if test_result[0]:
            print("API连接正常，开始性能测试...")
            main()
        else:
            print("API连接失败，请检查服务是否正常运行")
    except Exception as e:
        print(f"初始化失败: {str(e)}")
        print("请确保 requests 库已安装: pip install requests")