#!/usr/bin/env python3
"""
性能测试 - 测试大文件上传和并发
"""
import os
import tempfile
import time
import threading
from file_upload_sdk import init_sdk, upload_file

def test_large_file_upload():
    """测试大文件上传性能"""
    
    print("🚀 大文件上传性能测试")
    print("=" * 50)
    
    server_url = "https://file-upload-server.ai4s.com.cn"
    
    try:
        # 初始化SDK
        init_sdk(
            base_url=server_url,
            user_id=555,
            user_name="Performance Test User"
        )
        
        # 创建100MB测试文件
        print("\n📁 创建100MB测试文件...")
        start_time = time.time()
        
        with tempfile.NamedTemporaryFile(delete=False) as f:
            # 写入100MB数据 (1MB * 100)
            chunk = b"0" * (1024 * 1024)  # 1MB
            for i in range(100):
                f.write(chunk)
                if i % 20 == 0:
                    print(f"   写入进度: {i}%")
            large_file = f.name
        
        file_size = os.path.getsize(large_file)
        create_time = time.time() - start_time
        print(f"✅ 文件创建完成: {file_size / 1024 / 1024:.1f}MB, 耗时: {create_time:.2f}秒")
        
        # 上传性能测试
        print("\n📤 开始上传测试...")
        
        def progress_callback(percent):
            if percent % 10 == 0:  # 每10%显示一次
                print(f"   上传进度: {percent:.1f}%")
        
        upload_start = time.time()
        result = upload_file(
            file_path=large_file,
            dataset_id="performance_test_dataset",
            progress_callback=progress_callback
        )
        upload_time = time.time() - upload_start
        
        if result.get('success'):
            speed = file_size / upload_time / 1024 / 1024  # MB/s
            print(f"✅ 上传成功!")
            print(f"   文件大小: {file_size / 1024 / 1024:.1f}MB")
            print(f"   上传时间: {upload_time:.2f}秒")
            print(f"   上传速度: {speed:.2f}MB/s")
            print(f"   文件ID: {result.get('file_id')}")
        else:
            print(f"❌ 上传失败: {result.get('error')}")
        
        # 清理
        os.unlink(large_file)
        
    except Exception as e:
        print(f"❌ 大文件测试失败: {e}")
        if 'large_file' in locals():
            try:
                os.unlink(large_file)
            except:
                pass


def test_concurrent_uploads():
    """测试并发上传"""
    
    print("\n\n🔀 并发上传测试")
    print("=" * 50)
    
    server_url = "https://file-upload-server.ai4s.com.cn"
    
    def upload_worker(worker_id):
        """工作线程函数"""
        try:
            # 每个线程独立初始化SDK
            init_sdk(
                base_url=server_url,
                user_id=400 + worker_id,
                user_name=f"Concurrent Test User {worker_id}"
            )
            
            # 创建测试文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
                f.write(f"Concurrent upload test - Worker {worker_id}\n" * 1000)
                test_file = f.name
            
            # 上传文件
            start_time = time.time()
            result = upload_file(
                file_path=test_file,
                dataset_id=f"concurrent_test_{worker_id}"
            )
            upload_time = time.time() - start_time
            
            if result.get('success'):
                print(f"✅ Worker {worker_id}: 上传成功 ({upload_time:.2f}秒)")
            else:
                print(f"❌ Worker {worker_id}: 上传失败 - {result.get('error')}")
            
            # 清理
            os.unlink(test_file)
            
        except Exception as e:
            print(f"❌ Worker {worker_id}: 异常 - {e}")
    
    try:
        # 启动3个并发上传线程
        print("启动3个并发上传线程...")
        threads = []
        start_time = time.time()
        
        for i in range(3):
            thread = threading.Thread(target=upload_worker, args=(i,))
            threads.append(thread)
            thread.start()
        
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        
        total_time = time.time() - start_time
        print(f"\n✅ 并发测试完成，总耗时: {total_time:.2f}秒")
        
    except Exception as e:
        print(f"❌ 并发测试失败: {e}")


if __name__ == "__main__":
    print("🚀 Large File Upload SDK - 性能测试")
    print("📅 测试时间:", time.strftime('%Y-%m-%d %H:%M:%S'))
    print("⚠️  此测试会创建大文件，请确保有足够空间")
    
    # 询问是否执行大文件测试
    response = input("\n是否执行100MB大文件测试? (y/n): ").lower().strip()
    if response == 'y':
        test_large_file_upload()
    else:
        print("跳过大文件测试")
    
    # 并发测试
    test_concurrent_uploads()
    
    print("\n" + "=" * 50)
    print("🎯 性能测试完成！")
