#!/usr/bin/env python3
"""
综合测试 deadline 和 cancel_key 功能

验证两种任务控制机制可以正常协同工作
"""

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

# 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.processor import BatchProcessor
from sandai.operator.exceptions import TaskDeadlineExceededError, TaskCancelledError


def test_combined_deadline_and_cancel():
    """测试 deadline 和 cancel_key 的组合使用"""
    print("🧪 测试 deadline 和 cancel_key 的组合使用")
    
    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
        mock_result = Mock()
        mock_result.id = "combined-task-123"
        mock_app.send_task.return_value = mock_result
        
        # 创建客户端，同时设置默认 timeout 和 cancel_key
        client = OperatorClient(
            "test-operator", "v1", 
            timeout=300,  # 5分钟超时
            cancel_key="user_session_789"
        )
        
        # 发送任务，同时指定 timeout 和 cancel_key
        start_time = time.time()
        result = client.async_call(
            input_artifacts=[{"uri": "test.mp4", "slot": "video"}],
            options={"quality": "high"},
            timeout=120,  # 2分钟超时
            cancel_key="custom_session_456"
        )
        
        # 验证任务数据包含两个字段
        call_args = mock_app.send_task.call_args
        task_data = call_args[1]['kwargs']
        
        # 验证 deadline 字段
        assert 'deadline' in task_data, "任务数据应该包含 deadline 字段"
        expected_deadline = start_time + 120
        actual_deadline = task_data['deadline']
        assert abs(actual_deadline - expected_deadline) < 1, "deadline 计算错误"
        
        # 验证 cancel_key 字段
        assert 'cancel_key' in task_data, "任务数据应该包含 cancel_key 字段"
        assert task_data['cancel_key'] == "custom_session_456", "cancel_key 设置错误"
        
        print(f"   ✅ 同时设置 deadline 和 cancel_key:")
        print(f"      Deadline: {actual_deadline:.2f}")
        print(f"      Cancel key: {task_data['cancel_key']}")
        
        return True


def test_processor_priority_handling():
    """测试 processor 对 deadline 和 cancel_key 的优先级处理"""
    print("\n🧪 测试 processor 优先级处理")
    
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    current_time = time.time()
    
    # 模拟 Redis 客户端
    mock_redis = Mock()
    processor._redis_client = mock_redis
    
    print("1. 测试 deadline 优先于 cancel_key:")
    
    # 创建同时过期和被取消的任务
    expired_and_cancelled_task = {
        "task_id": "expired-cancelled-task",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "deadline": current_time - 10,  # 已过期
        "cancel_key": "cancelled_session"
    }
    
    # 设置 Redis 返回 key 存在
    mock_redis.exists.return_value = True
    
    # deadline 检查应该优先
    should_skip_deadline = processor._should_skip_due_to_deadline(expired_and_cancelled_task)
    should_skip_cancel = processor._should_skip_due_to_cancel(expired_and_cancelled_task)
    
    assert should_skip_deadline == True, "过期任务应该被 deadline 检查捕获"
    assert should_skip_cancel == True, "被取消任务应该被 cancel 检查捕获"
    
    print("   ✅ deadline 和 cancel 检查都能正确识别问题")
    
    print("\n2. 测试只有 cancel_key 的情况:")
    
    # 创建只被取消的任务
    only_cancelled_task = {
        "task_id": "only-cancelled-task",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "deadline": current_time + 300,  # 未过期
        "cancel_key": "cancelled_session_2"
    }
    
    should_skip_deadline = processor._should_skip_due_to_deadline(only_cancelled_task)
    should_skip_cancel = processor._should_skip_due_to_cancel(only_cancelled_task)
    
    assert should_skip_deadline == False, "未过期任务不应该被 deadline 检查捕获"
    assert should_skip_cancel == True, "被取消任务应该被 cancel 检查捕获"
    
    print("   ✅ 只有 cancel 的情况正确处理")
    
    print("\n3. 测试只有 deadline 的情况:")
    
    # 创建只过期的任务
    only_expired_task = {
        "task_id": "only-expired-task",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "deadline": current_time - 5,  # 已过期
        # 没有 cancel_key
    }
    
    should_skip_deadline = processor._should_skip_due_to_deadline(only_expired_task)
    should_skip_cancel = processor._should_skip_due_to_cancel(only_expired_task)
    
    assert should_skip_deadline == True, "过期任务应该被 deadline 检查捕获"
    assert should_skip_cancel == False, "无 cancel_key 任务不应该被 cancel 检查捕获"
    
    print("   ✅ 只有 deadline 的情况正确处理")
    
    return True


def test_error_output_types():
    """测试不同类型的错误输出"""
    print("\n🧪 测试不同类型的错误输出")
    
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    current_time = time.time()
    
    # 测试 deadline 错误输出
    expired_task = {
        "task_id": "expired-task-123",
        "deadline": current_time - 30,
        "cancel_key": "some_key"
    }
    
    deadline_error = processor._create_deadline_error_output(expired_task)
    
    assert deadline_error["error"]["error_type"] == "TaskDeadlineExceededError"
    assert "deadline exceeded" in deadline_error["error"]["error_message"]
    assert deadline_error["error"]["error_details"]["exceeded_by"] > 0
    
    print("   ✅ deadline 错误输出格式正确")
    
    # 测试 cancel 错误输出
    cancelled_task = {
        "task_id": "cancelled-task-456",
        "cancel_key": "user_session_123"
    }
    
    cancel_error = processor._create_cancel_error_output(cancelled_task)
    
    assert cancel_error["error"]["error_type"] == "TaskCancelledError"
    assert "cancelled via cancel_key" in cancel_error["error"]["error_message"]
    assert cancel_error["error"]["error_details"]["cancel_key"] == "user_session_123"
    
    print("   ✅ cancel 错误输出格式正确")
    
    # 验证两种错误输出都有共同的标记
    assert deadline_error["__keep_raw_output__"] == True
    assert cancel_error["__keep_raw_output__"] == True
    assert deadline_error["error"]["error_details"]["skipped_before_processing"] == True
    assert cancel_error["error"]["error_details"]["skipped_before_processing"] == True
    
    print("   ✅ 错误输出共同标记正确")
    
    return True


def test_performance_optimization():
    """测试性能优化特性"""
    print("\n🧪 测试性能优化特性")
    
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    # 模拟 Redis 客户端
    mock_redis = Mock()
    processor._redis_client = mock_redis
    
    # 测试已知被取消的 key 不会重复检查 Redis
    cancelled_task = {
        "task_id": "test-task",
        "cancel_key": "known_cancelled_key"
    }
    
    # 第一次检查，Redis 返回存在
    mock_redis.exists.return_value = True
    should_skip_1 = processor._should_skip_due_to_cancel(cancelled_task)
    
    # 验证第一次调用了 Redis
    assert mock_redis.exists.called, "第一次应该调用 Redis"
    assert should_skip_1 == True, "第一次应该跳过"
    
    # 重置 mock
    mock_redis.reset_mock()
    
    # 第二次检查同样的 key
    should_skip_2 = processor._should_skip_due_to_cancel(cancelled_task)
    
    # 验证第二次没有调用 Redis
    assert not mock_redis.exists.called, "第二次不应该调用 Redis"
    assert should_skip_2 == True, "第二次也应该跳过"
    
    print("   ✅ 已知被取消的 key 不会重复检查 Redis")
    
    # 验证 cancelled_keys 集合包含了这个 key
    assert "known_cancelled_key" in processor._cancelled_keys, "key 应该被添加到已知集合中"
    
    print("   ✅ cancelled_keys 集合正确维护")
    
    return True


def run_all_tests():
    """运行所有综合测试"""
    print("🚀 开始综合测试 deadline 和 cancel_key 功能")
    print("=" * 70)
    
    tests = [
        test_combined_deadline_and_cancel,
        test_processor_priority_handling,
        test_error_output_types,
        test_performance_optimization,
    ]
    
    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" + "=" * 70)
    print(f"📊 测试结果: {passed} 通过, {failed} 失败")
    
    if failed == 0:
        print("🎉 所有综合测试通过!")
        print("\n📋 综合功能总结:")
        print("   🎯 Deadline 机制: 基于时间的任务超时控制")
        print("   🔑 Cancel Key 机制: 基于 Redis 的任务取消控制")
        print("   ⚡ 早期检查: 两种机制都在处理前检查，节省资源")
        print("   🛡️ 错误处理: 不同的错误类型，便于客户端区分处理")
        print("   🚀 性能优化: 缓存已知状态，避免重复检查")
        print("\n✨ 使用场景:")
        print("   - Deadline: 实时处理、批量作业、资源限制场景")
        print("   - Cancel Key: 用户取消、会话过期、业务逻辑取消")
        return True
    else:
        print("💥 部分测试失败，请检查实现")
        return False


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