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

验证 client 发送任务时会设置 cancel_key，processor 会检查 Redis 并处理被取消的任务
"""

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 TaskCancelledError, TaskExecutionError


def test_client_cancel_key():
    """测试客户端 cancel_key 功能"""
    print("🧪 测试客户端 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
        
        # 模拟 send_task 返回值
        mock_result = Mock()
        mock_result.id = "test-task-123"
        mock_app.send_task.return_value = mock_result
        
        print("1. 测试默认 cancel_key:")
        
        # 创建带默认 cancel_key 的客户端
        client = OperatorClient("test-operator", "v1", cancel_key="session_456")
        
        # 发送任务
        result = client.async_call(
            input_artifacts=[{"uri": "test.mp4", "slot": "video"}],
            options={"test": True}
        )
        
        # 验证 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 参数
        
        # 验证 cancel_key 字段存在
        assert 'cancel_key' in task_data, "任务数据应该包含 cancel_key 字段"
        assert task_data['cancel_key'] == "session_456", f"cancel_key 应该是 'session_456'，实际是 {task_data['cancel_key']}"
        
        print(f"   ✅ 默认 cancel_key 正确设置: {task_data['cancel_key']}")
        
        print("\n2. 测试自定义 cancel_key:")
        
        # 发送带自定义 cancel_key 的任务
        result2 = client.async_call(
            input_artifacts=[{"uri": "test2.mp4", "slot": "video"}],
            options={"test": True},
            cancel_key="custom_789"
        )
        
        # 获取调用参数
        call_args = mock_app.send_task.call_args
        task_data = call_args[1]['kwargs']
        
        assert task_data['cancel_key'] == "custom_789", f"cancel_key 应该是 'custom_789'，实际是 {task_data['cancel_key']}"
        
        print(f"   ✅ 自定义 cancel_key 正确设置: {task_data['cancel_key']}")
        
        print("\n3. 测试无 cancel_key:")
        
        # 创建无默认 cancel_key 的客户端
        client_no_cancel = OperatorClient("test-operator", "v1")
        
        result3 = client_no_cancel.async_call(
            input_artifacts=[{"uri": "test3.mp4", "slot": "video"}],
            options={"test": True}
        )
        
        # 获取调用参数
        call_args = mock_app.send_task.call_args
        task_data = call_args[1]['kwargs']
        
        assert 'cancel_key' not in task_data, "任务数据不应该包含 cancel_key 字段"
        
        print("   ✅ 无 cancel_key 时正确不设置字段")
        
        print("✅ 客户端 cancel_key 功能测试通过")
        return True


def test_client_cancel_error_handling():
    """测试客户端 cancel 错误处理"""
    print("\n🧪 测试客户端 cancel 错误处理")
    
    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
        
        # 模拟 cancel 错误的返回结果
        cancel_error_result = {
            'status': 'error',
            'error': {
                'error_type': 'TaskCancelledError',
                'error_message': 'Task was cancelled via cancel_key: session_123',
                'error_details': {
                    'cancel_key': 'session_123',
                    'skipped_before_processing': True
                }
            }
        }
        
        mock_result.get.return_value = cancel_error_result
        
        # 创建客户端
        client = OperatorClient("test-operator", "v1")
        
        # 测试同步调用应该抛出 TaskCancelledError
        try:
            client.sync(
                input_artifacts=[{"uri": "test.mp4", "slot": "video"}],
                options={"test": True},
                cancel_key="session_123"
            )
            assert False, "应该抛出 TaskCancelledError"
        except TaskCancelledError as e:
            print(f"   ✅ 正确捕获 TaskCancelledError: {e}")
            assert e.task_id == "test-task-456"
            assert e.cancel_key == "session_123"
        except Exception as e:
            assert False, f"应该抛出 TaskCancelledError，但抛出了 {type(e).__name__}: {e}"
        
        print("✅ 客户端 cancel 错误处理测试通过")
        return True


def test_processor_cancel_check():
    """测试 processor cancel 检查逻辑"""
    print("\n🧪 测试 processor cancel 检查逻辑")
    
    from sandai.operator.processor import BatchProcessor
    
    # 创建 processor
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    print("1. 测试 cancel_key 检查逻辑:")
    
    # 创建有 cancel_key 的任务输入
    task_with_cancel = {
        "task_id": "cancel-task-123",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True},
        "cancel_key": "session_456"
    }
    
    # 创建无 cancel_key 的任务输入
    task_no_cancel = {
        "task_id": "normal-task-456",
        "artifacts": [{"uri": "test.mp4", "slot": "video"}],
        "options": {"test": True}
    }
    
    # 模拟 Redis 客户端
    mock_redis = Mock()
    processor._redis_client = mock_redis
    
    # 测试 cancel_key 不存在的情况
    mock_redis.exists.return_value = False
    should_skip = processor._should_skip_due_to_cancel(task_with_cancel)
    assert should_skip == False, "cancel_key 不存在时不应该跳过"
    print("   ✅ cancel_key 不存在时正确不跳过")
    
    # 测试 cancel_key 存在的情况
    mock_redis.exists.return_value = True
    should_skip = processor._should_skip_due_to_cancel(task_with_cancel)
    assert should_skip == True, "cancel_key 存在时应该跳过"
    print("   ✅ cancel_key 存在时正确跳过")
    
    # 测试已知被取消的 key
    should_skip = processor._should_skip_due_to_cancel(task_with_cancel)
    assert should_skip == True, "已知被取消的 key 应该直接跳过"
    print("   ✅ 已知被取消的 key 正确直接跳过")
    
    # 测试无 cancel_key 的任务
    should_skip = processor._should_skip_due_to_cancel(task_no_cancel)
    assert should_skip == False, "无 cancel_key 的任务不应该跳过"
    print("   ✅ 无 cancel_key 的任务正确不跳过")
    
    print("\n2. 测试 cancel 错误输出创建:")
    
    error_output = processor._create_cancel_error_output(task_with_cancel)
    
    assert error_output["task_id"] == "cancel-task-123"
    assert error_output["status"] == "error"
    assert error_output["error"]["error_type"] == "TaskCancelledError"
    assert "cancelled via cancel_key" in error_output["error"]["error_message"]
    assert error_output["error"]["error_details"]["cancel_key"] == "session_456"
    assert error_output["error"]["error_details"]["skipped_before_processing"] == True
    assert error_output["__keep_raw_output__"] == True
    
    print("   ✅ cancel 错误输出创建正确")
    print(f"      错误类型: {error_output['error']['error_type']}")
    print(f"      Cancel key: {error_output['error']['error_details']['cancel_key']}")
    print(f"      处理前跳过: {error_output['error']['error_details']['skipped_before_processing']}")
    
    print("✅ processor cancel 检查逻辑测试通过")
    return True


def test_redis_client_setup():
    """测试 Redis 客户端设置"""
    print("\n🧪 测试 Redis 客户端设置")
    
    from sandai.operator.processor import BatchProcessor
    
    # 创建 processor
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    
    # 模拟 Celery channel
    mock_channel = Mock()
    mock_celery_app = Mock()
    mock_celery_app.conf.broker_url = "redis://localhost:6379/1"
    mock_channel.celery_app = mock_celery_app
    
    # 模拟 Redis 客户端
    with patch('redis.Redis') as mock_redis_class:
        mock_redis_instance = Mock()
        mock_redis_class.return_value = mock_redis_instance
        
        # 测试 Redis 客户端设置
        processor._setup_redis_client(mock_channel)
        
        # 验证 Redis 客户端被创建
        mock_redis_class.assert_called_once_with(
            host='localhost',
            port=6379,
            db=1,
            decode_responses=True
        )
        
        # 验证 ping 被调用
        mock_redis_instance.ping.assert_called_once()
        
        # 验证 Redis 客户端被设置
        assert processor._redis_client == mock_redis_instance
        
        print("   ✅ Redis 客户端设置正确")
        print(f"      Host: localhost")
        print(f"      Port: 6379")
        print(f"      DB: 1")
    
    print("✅ Redis 客户端设置测试通过")
    return True


def run_all_tests():
    """运行所有测试"""
    print("🚀 开始测试 cancel_key 功能")
    print("=" * 60)
    
    tests = [
        test_client_cancel_key,
        test_client_cancel_error_handling,
        test_processor_cancel_check,
        test_redis_client_setup,
    ]
    
    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("🎉 所有 cancel_key 功能测试通过!")
        print("\n📋 功能总结:")
        print("   ✅ 客户端发送任务时可以设置 cancel_key")
        print("   ✅ processor 从 Celery app 获取 Redis 客户端")
        print("   ✅ processor 检查 Redis 中的 cancel_key")
        print("   ✅ 维护已知被取消的 key 集合，避免重复检查")
        print("   ✅ 被取消的任务在处理前就跳过，节省资源")
        print("   ✅ 客户端正确处理任务取消异常")
        return True
    else:
        print("💥 部分测试失败，请检查实现")
        return False


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