#!/usr/bin/env python3
"""
异步批处理器测试

测试 Sandai Operator SDK 的 BatchProcessor 异步批处理功能
"""

import json
import time
import sys
import os
from pathlib import Path

from sandai.operator import BatchProcessor, TaskInput, TaskOutput

# 模拟的处理选项和结果类型
from pydantic import BaseModel

class TestOptions(BaseModel):
    delay: float = 0.1

class TestResults(BaseModel):
    processed_at: str
    batch_size: int

def test_async_batch_processing():
    """测试异步批处理功能"""
    
    # 清理测试文件
    test_files = ["async_test_input.jsonl", "async_test_output.jsonl"]
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    # 创建处理器实例
    processor = BatchProcessor(
        name="test-async-processor",
        version="1.0.0",
        max_concurrency=2,  # 最大并发批处理器数量
        max_batch_size=3,   # 每个批次最大任务数量
        max_prefetch_size=10, # 输入队列最大大小
        ignore_nongil_check=True
    )

    @processor.on_batch(max_concurrency=2, max_batch_size=3, max_prefetch_size=10)
    def process_batch(batch_inputs: list[TaskInput[TestOptions]], operator_config: dict, context):
        """模拟批处理函数"""
        print(f"处理批次: {len(batch_inputs)} 个任务")
        
        # 模拟计算密集型操作
        max_delay = max(task.options.delay for task in batch_inputs)
        time.sleep(max_delay)  # 模拟处理时间
        
        # 生成结果
        for task_input in batch_inputs:
            yield TaskOutput(
                task_id=task_input.task_id,  # 使用 task_id
                results=TestResults(
                    processed_at=time.strftime("%Y-%m-%d %H:%M:%S"),
                    batch_size=len(batch_inputs)
                ),
                artifacts=[]
            )

    # 创建测试输入文件
    test_inputs = []
    for i in range(8):  # 创建8个测试任务
        test_inputs.append({
            "task_id": f"async-task-{i:03d}",  # 使用 task_id 而不是 id
            "options": {"delay": 0.1 + (i % 3) * 0.05},  # 不同的处理时间
            "artifacts": []
        })
    
    with open("async_test_input.jsonl", "w") as f:
        for input_data in test_inputs:
            f.write(json.dumps(input_data) + "\n")
    
    print(f"创建了 {len(test_inputs)} 个测试任务")
    
    # 配置处理器
    processor.config = {
        "channel.file.input_path": "async_test_input.jsonl",
        "channel.file.output_path": "async_test_output.jsonl",
        "operator.config": {}
    }
    
    print("启动异步批处理器...")
    print(f"配置: max_batch_size={processor.max_batch_size}, max_concurrency={processor.max_concurrency}")
    
    start_time = time.time()
    
    try:
        # 运行处理器（设置超时）
        import signal
        def timeout_handler(signum, frame):
            raise TimeoutError("处理超时")
        
        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(15)  # 15秒超时
        
        processor.run()
        
    except (KeyboardInterrupt, TimeoutError):
        print("\n处理完成或超时")
    finally:
        signal.alarm(0)  # 取消超时
    
    end_time = time.time()
    processing_time = end_time - start_time
    print(f"\n处理耗时: {processing_time:.2f} 秒")
    
    # 检查输出结果
    success = False
    if Path("async_test_output.jsonl").exists():
        with open("async_test_output.jsonl", "r") as f:
            results = [json.loads(line) for line in f]
        
        print(f"生成了 {len(results)} 个输出结果")
        
        # 分析批次分布
        batch_sizes = {}
        for result in results:
            batch_size = result.get("results", {}).get("batch_size", 1)
            batch_sizes[batch_size] = batch_sizes.get(batch_size, 0) + 1
        
        print("批次大小分布:")
        for size, count in sorted(batch_sizes.items()):
            batches = count // size
            print(f"  批次大小 {size}: {count} 个任务, {batches} 个批次")
        
        # 计算理论性能提升
        if len(results) > 0:
            avg_task_time = 0.125  # 平均任务处理时间
            serial_time = len(test_inputs) * avg_task_time
            avg_batch_size = sum(size * count for size, count in batch_sizes.items()) / len(results)
            
            print(f"\n性能分析:")
            print(f"  串行处理预估时间: {serial_time:.2f} 秒")
            print(f"  实际处理时间: {processing_time:.2f} 秒")
            if processing_time > 0:
                speedup = serial_time / processing_time
                print(f"  加速比: {speedup:.2f}x")
        else:
            print(f"\n❌ 没有处理任何任务")
        
        success = len(results) >= len(test_inputs) * 0.8  # 至少处理80%的任务
    else:
        print("❌ 没有找到输出文件")
    
    # 清理测试文件
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    return success

def main():
    """运行异步批处理测试"""
    print("异步批处理器测试")
    print("=" * 50)
    
    try:
        success = test_async_batch_processing()
        
        if success:
            print("\n✅ 异步批处理测试通过！")
            return 0
        else:
            print("\n❌ 异步批处理测试失败！")
            return 1
            
    except Exception as e:
        print(f"\n❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    exit(main())