#!/usr/bin/env python3
"""
作业模式测试

测试 Sandai Operator SDK 的作业模式功能
"""

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

from sandai.operator import BatchProcessor, TaskInput, TaskOutput
from pydantic import BaseModel

# 测试用的选项和结果类型
class JobTestOptions(BaseModel):
    delay: float = 0.1
    message: str = "test"

class JobTestResults(BaseModel):
    processed_at: str
    batch_size: int
    message: str

def test_job_mode():
    """测试作业模式"""
    
    # 清理测试文件
    test_files = ["job_mode_input.jsonl", "job_mode_output.jsonl"]
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    # 创建处理器实例
    processor = BatchProcessor(
        name="test-job-mode",
        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[JobTestOptions]], operator_config: dict, context):
        """批处理函数"""
        print(f"处理批次: {len(batch_inputs)} 个任务")
        
        # 模拟处理时间
        time.sleep(0.1)
        
        # 生成结果
        for task_input in batch_inputs:
            yield TaskOutput(
                task_id=task_input.task_id,
                results=JobTestResults(
                    processed_at=time.strftime("%Y-%m-%d %H:%M:%S"),
                    batch_size=len(batch_inputs),
                    message=f"Processed: {task_input.options.message}"
                ),
                artifacts=[]
            )

    # 创建测试输入文件
    test_inputs = []
    for i in range(7):  # 创建7个测试任务
        test_inputs.append({
            "task_id": f"job-mode-{i:03d}",
            "options": {"delay": 0.1, "message": f"task_{i}"},
            "artifacts": []
        })
    
    with open("job_mode_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": "job_mode_input.jsonl",
        "channel.file.output_path": "job_mode_output.jsonl",
        "operator.config": {}
    }
    
    print("启动作业模式...")
    start_time = time.time()
    
    try:
        # 运行作业模式
        processor.run(job_mode=True)
        
    except Exception as e:
        print(f"处理异常: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    end_time = time.time()
    processing_time = end_time - start_time
    print(f"作业模式耗时: {processing_time:.2f} 秒")
    
    # 检查输出结果
    success = False
    if Path("job_mode_output.jsonl").exists():
        with open("job_mode_output.jsonl", "r") as f:
            results = [json.loads(line) for line in f]
        
        print(f"生成了 {len(results)} 个输出结果")
        
        # 验证结果
        if len(results) == len(test_inputs):
            print("✅ 所有任务都被处理了")
            
            # 检查批次分布
            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} 个批次")
            
            success = True
        else:
            print(f"❌ 任务数量不匹配: 期望 {len(test_inputs)}, 实际 {len(results)}")
    else:
        print("❌ 没有找到输出文件")
    
    # 清理测试文件
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    return success

def test_empty_input_job_mode():
    """测试空输入的作业模式"""
    
    # 清理测试文件
    test_files = ["empty_job_input.jsonl", "empty_job_output.jsonl"]
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    # 创建空的输入文件
    with open("empty_job_input.jsonl", "w") as f:
        pass  # 空文件
    
    # 创建处理器
    processor = BatchProcessor(
        name="test-empty-job",
        version="1.0.0",
        max_concurrency=1,
        max_batch_size=2,
        ignore_nongil_check=True
    )

    @processor.on_batch()
    def process_batch(batch_inputs, operator_config, context):
        for task_input in batch_inputs:
            yield TaskOutput(
                task_id=task_input.task_id,
                results={"processed": True},
                artifacts=[]
            )

    processor.config = {
        "channel.file.input_path": "empty_job_input.jsonl",
        "channel.file.output_path": "empty_job_output.jsonl",
        "operator.config": {}
    }
    
    print("\n测试空输入的作业模式...")
    start_time = time.time()
    
    try:
        processor.run(job_mode=True)
        end_time = time.time()
        processing_time = end_time - start_time
        
        print(f"空输入作业模式耗时: {processing_time:.2f} 秒")
        
        # 应该很快退出，没有输出文件或输出文件为空
        if not Path("empty_job_output.jsonl").exists():
            print("✅ 空输入正确处理 - 没有生成输出文件")
            success = True
        else:
            with open("empty_job_output.jsonl", "r") as f:
                content = f.read().strip()
            if not content:
                print("✅ 空输入正确处理 - 输出文件为空")
                success = True
            else:
                print("❌ 空输入处理异常 - 输出文件不为空")
                success = False
        
    except Exception as e:
        print(f"❌ 空输入处理异常: {e}")
        success = False
    
    # 清理测试文件
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()
    
    return success

def main():
    """运行作业模式测试"""
    print("作业模式测试")
    print("=" * 50)
    
    success1 = test_job_mode()
    success2 = test_empty_input_job_mode()
    
    if success1 and success2:
        print("\n✅ 所有作业模式测试通过！")
        return 0
    else:
        print("\n❌ 作业模式测试失败！")
        return 1

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