#!/usr/bin/env python3
"""
测试 deadline 功能

验证 client 发送任务时会设置 deadline，processor 会检查并处理超时任务
"""

import sys
import time
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock

# Add the parent directory to the path to import sandai
sys.path.insert(0, str(Path(__file__).parent.parent))

from sandai.operator_client import OperatorClient
from sandai.operator.exceptions import TaskDeadlineExceededError, TaskExecutionError


def test_client_deadline_calculation():
    """测试客户端 deadline 计算"""
    print("🧪 测试客户端 deadline 计算")
    
    with patch('sandai.operator_client.Celery') as mock_celery, \
         patch('sandai.operator.env.get_celery_config') as mock_get_config, \
         patch('sandai.operator.env.get_default_celery_config') as mock_default_config:
        
        # 模拟配置
        mock_config = Mock()
        mock_config.broker_url = "redis://localhost:6379/0"
        mock_config.result_backend = "redis://localhost:6379/0"
        
        mock_get_config.return_value = mock_config
        mock_default_config.return_value = mock_config
        
        # 模拟 Celery 应用
        mock_app = Mock()
        mock_celery.return_value = mock_app
        
        # 模拟 send_task 返回值
        mock_result = Mock()
        mock_result.id = "test-task-123"
        mock_app.send_task.return_value = mock_result
        
        # 创建客户端
        client = OperatorClient("test-operator", "v1", timeout=300)  # 5分钟超时
        
        # 记录发送时间
        start_time = time.time()
        
        # 发送任务
        result = client.async_call(
            input_artifacts=[{"uri": "test.mp4", "slot": "video"}],
            options={"test": True},
            timeout=120  # 2分钟超时
        )
        
        # 验证 send_task 被调用
        assert mock_app.send_task.called, "send_task 应该被调用"
        
        # 获取调用参数
        call_args = mock_app.send_task.call_args
        task_data = call_args[1]['kwargs']  # kwargs 参数
        
        # 验证 deadline 字段存在
        assert 'deadline' in task_data, "任务数据应该包含 deadline 字段"
        
        # 验证 deadline 值合理（应该是当前时间 + 120秒左右）
        expected_deadline = start_time + 120
        actual_deadline = task_data['deadline']
        
        # 允许1秒的误差
        assert abs(actual_deadline - expected_deadline) < 1, \
            f"deadline 计算错误: 期望 {expected_deadline}, 实际 {actual_deadline}"
        
        print(f"✅ deadline 计算正确: {actual_deadline:.2f} (期望: {expected_deadline:.2f})")
        
        # 验证其他字段
        assert 'artifacts' in task_data, "任务数据应该包含 artifacts 字段"
        assert 'options' in task_data, "任务数据应该包含 options 字段"
        
        print("✅ 客户端 deadline 计算测试通过")
        return True


def test_client_deadline_error_handling():
    """测试客户端 deadline 错误处理"""
    print("\n🧪 测试客户端 deadline 错误处理")
    
    with patch('sandai.operator_client.Celery') as mock_celery, \
         patch('sandai.operator.env.get_celery_config') as mock_get_config, \
         patch('sandai.operator.env.get_default_celery_config') as mock_default_config:
        
        # 模拟配置
        mock_config = Mock()
        mock_config.broker_url = "redis://localhost:6379/0"
        mock_config.result_backend = "redis://localhost:6379/0"
        
        mock_get_config.return_value = mock_config
        mock_default_config.return_value = mock_config
        
        # 模拟 Celery 应用
        mock_app = Mock()
        mock_celery.return_value = mock_app
        
        # 模拟 AsyncResult
        mock_result = Mock()
        mock_result.id = "test-task-456"
        mock_app.send_task.return_value = mock_result
        
        # 模拟 deadline 超时的返回结果
        deadline_error_result = {
            'status': 'error',
            'error': {
                'error_type': 'TaskDeadlineExceededError',
                'error_message': 'Task deadline exceeded: 1234567890.0 < 1234567900.0',
                'error_details': {
                    'deadline': 1234567890.0,
                    'current_time': 1234567900.0,
                    'exceeded_by': 10.0
                }
            }
        }
        
        mock_result.get.return_value = deadline_error_result
        
        # 创建客户端
        client = OperatorClient("test-operator", "v1")
        
        # 测试同步调用应该抛出 TaskDeadlineExceededError
        try:
            client.sync(
                input_artifacts=[{"uri": "test.mp4", "slot": "video"}],
                options={"test": True},
                timeout=60
            )
            assert False, "应该抛出 TaskDeadlineExceededError"
        except TaskDeadlineExceededError as e:
            print(f"✅ 正确捕获 TaskDeadlineExceededError: {e}")
            assert e.task_id == "test-task-456"
            assert e.deadline == 1234567890.0
            assert e.current_time == 1234567900.0
        except Exception as e:
            assert False, f"应该抛出 TaskDeadlineExceededError，但抛出了 {type(e).__name__}: {e}"
        
        print("✅ 客户端 deadline 错误处理测试通过")
        return True


def test_processor_deadline_check():
    """测试 processor deadline 检查逻辑"""
    print("\n🧪 测试 processor deadline 检查逻辑")
    
    from sandai.operator.processor import BatchProcessor
    
    # 创建 processor
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    current_time = time.time()
    
    # 测试 _should_skip_due_to_deadline 方法
    
    # 创建一个已过期的任务输入
    expired_input = {
        "task_id": "expired-task-123",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "deadline": current_time - 10  # 10秒前就过期了
    }
    
    # 创建一个未过期的任务输入
    valid_input = {
        "task_id": "valid-task-456",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "deadline": current_time + 300  # 5分钟后过期
    }
    
    # 创建一个没有 deadline 的任务输入
    no_deadline_input = {
        "task_id": "no-deadline-task-789",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True}
    }
    
    # 测试 deadline 检查逻辑
    assert processor._should_skip_due_to_deadline(expired_input) == True, "过期任务应该被跳过"
    print("✅ 过期任务正确识别为应跳过")
    
    assert processor._should_skip_due_to_deadline(valid_input) == False, "有效任务不应该被跳过"
    print("✅ 有效任务正确识别为不跳过")
    
    assert processor._should_skip_due_to_deadline(no_deadline_input) == False, "无 deadline 任务不应该被跳过"
    print("✅ 无 deadline 任务正确识别为不跳过")
    
    # 测试 _create_deadline_error_output 方法
    error_output = processor._create_deadline_error_output(expired_input)
    
    assert error_output["task_id"] == "expired-task-123"
    assert error_output["status"] == "error"
    assert error_output["error"]["error_type"] == "TaskDeadlineExceededError"
    assert "deadline exceeded before processing" in error_output["error"]["error_message"]
    assert error_output["error"]["error_details"]["skipped_before_processing"] == True
    assert error_output["error"]["error_details"]["exceeded_by"] > 0
    assert error_output["__keep_raw_output__"] == True
    
    print("✅ deadline 错误输出创建正确")
    print(f"   错误类型: {error_output['error']['error_type']}")
    print(f"   跳过标记: {error_output['error']['error_details']['skipped_before_processing']}")
    print(f"   超时时间: {error_output['error']['error_details']['exceeded_by']:.2f} 秒")
    
    print("✅ processor deadline 检查逻辑测试通过")
    return True


def test_error_output_creation():
    """测试错误输出创建"""
    print("\n🧪 测试错误输出创建")
    
    from sandai.operator.processor import BatchProcessor
    
    # 创建 processor
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    # 测试创建 deadline 超时错误输出
    current_time = time.time()
    deadline = current_time - 10
    
    error_output = processor.create_error_output(
        task_id="test-task-123",
        error_type="TaskDeadlineExceededError",
        error_message=f"Task deadline exceeded: {deadline} < {current_time}",
        error_details={
            "deadline": deadline,
            "current_time": current_time,
            "exceeded_by": current_time - deadline
        }
    )
    
    # 验证错误输出
    assert error_output.task_id == "test-task-123"
    assert error_output.status == "error"
    assert error_output.error is not None
    assert error_output.error.error_type == "TaskDeadlineExceededError"
    assert "deadline exceeded" in error_output.error.error_message.lower()
    assert error_output.error.error_details["deadline"] == deadline
    assert error_output.error.error_details["current_time"] == current_time
    assert error_output.error.error_details["exceeded_by"] > 0
    
    print("✅ deadline 超时错误输出创建正确")
    print(f"   错误类型: {error_output.error.error_type}")
    print(f"   错误消息: {error_output.error.error_message}")
    print(f"   超时时间: {error_output.error.error_details['exceeded_by']:.2f} 秒")
    
    print("✅ 错误输出创建测试通过")
    return True


def run_all_tests():
    """运行所有测试"""
    print("🚀 开始测试 deadline 功能")
    print("=" * 60)
    
    tests = [
        test_client_deadline_calculation,
        test_client_deadline_error_handling,
        test_processor_deadline_check,
        test_error_output_creation,
    ]
    
    passed = 0
    failed = 0
    
    for test in tests:
        try:
            if test():
                passed += 1
            else:
                failed += 1
        except Exception as e:
            print(f"❌ 测试失败: {test.__name__}")
            print(f"   错误: {e}")
            import traceback
            traceback.print_exc()
            failed += 1
    
    print("\n" + "=" * 60)
    print(f"📊 测试结果: {passed} 通过, {failed} 失败")
    
    if failed == 0:
        print("🎉 所有 deadline 功能测试通过!")
        print("\n📋 功能总结:")
        print("   ✅ 客户端发送任务时自动计算 deadline")
        print("   ✅ processor 在接收任务时立即检查 deadline（避免资源浪费）")
        print("   ✅ deadline 超时任务在下载输入文件前就被跳过")
        print("   ✅ 正确创建 deadline 超时错误输出")
        print("   ✅ 客户端正确处理 deadline 超时异常")
        return True
    else:
        print("💥 部分测试失败，请检查实现")
        return False


if __name__ == '__main__':
    success = run_all_tests()
    sys.exit(0 if success else 1)