#!/usr/bin/env python3
"""
测试简化的 Celery Channel 使用方式

验证客户端可以使用 app.send_task(artifacts, options) 的简化方式，
而处理器会自动使用 Celery 生成的任务 ID。
"""

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

# 添加 sandai 包到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from sandai.operator.channel import CeleryChannel


def test_simplified_task_format():
    """测试简化的任务格式处理"""
    print("Testing simplified task format...")
    
    # 模拟 Celery 消息
    mock_message = Mock()
    mock_message.properties = {'correlation_id': 'celery_task_12345'}
    mock_message.id = 'celery_task_12345'
    mock_message.ack = Mock()
    mock_message.reject = Mock()
    
    # 模拟客户端发送的简化任务数据（不包含 task_id）
    client_task_data = {
        "artifacts": [
            {
                "uri": "file://test.txt",
                "slot": "input",
                "name": "test.txt"
            }
        ],
        "options": {
            "param1": "value1",
            "param2": "value2"
        }
    }
    
    # 模拟 CeleryChannel 的 process_task 逻辑
    celery_task_id = mock_message.properties.get('correlation_id')
    
    # 复制任务数据并添加 task_id
    processed_task_data = client_task_data.copy()
    processed_task_data['task_id'] = celery_task_id
    
    # 验证处理后的任务数据
    assert processed_task_data['task_id'] == 'celery_task_12345'
    assert processed_task_data['artifacts'] == client_task_data['artifacts']
    assert processed_task_data['options'] == client_task_data['options']
    
    print("✓ Simplified task format processing test passed")


def test_celery_id_extraction():
    """测试 Celery ID 提取逻辑"""
    print("Testing Celery ID extraction...")
    
    # 测试从 correlation_id 提取
    mock_message1 = Mock()
    mock_message1.properties = {'correlation_id': 'test_id_123'}
    mock_message1.id = 'fallback_id'
    
    celery_id = mock_message1.properties.get('correlation_id') or getattr(mock_message1, 'id', None)
    assert celery_id == 'test_id_123'
    
    # 测试从 message.id 提取（当 correlation_id 不存在时）
    mock_message2 = Mock()
    mock_message2.properties = {}
    mock_message2.id = 'message_id_456'
    
    celery_id = mock_message2.properties.get('correlation_id') or getattr(mock_message2, 'id', None)
    assert celery_id == 'message_id_456'
    
    print("✓ Celery ID extraction test passed")


def test_task_data_structure():
    """测试任务数据结构"""
    print("Testing task data structure...")
    
    # 客户端发送的数据
    client_data = {
        "artifacts": [],
        "options": {
            "message": "test",
            "delay": 1
        }
    }
    
    # 处理器接收到的数据（添加了 task_id）
    celery_task_id = "auto_generated_id_789"
    processor_data = client_data.copy()
    processor_data['task_id'] = celery_task_id
    
    # 验证数据结构
    expected_keys = {'task_id', 'artifacts', 'options'}
    assert set(processor_data.keys()) == expected_keys
    assert processor_data['task_id'] == celery_task_id
    
    print("✓ Task data structure test passed")


def main():
    """主函数"""
    print("Simplified Celery Channel Usage Test")
    print("=" * 50)
    
    try:
        test_simplified_task_format()
        test_celery_id_extraction()
        test_task_data_structure()
        
        print("\n✅ All simplified usage tests passed!")
        print("\n📝 验证的功能:")
        print("   ✓ 客户端可以发送不包含 task_id 的任务")
        print("   ✓ CeleryChannel 自动提取 Celery 任务 ID")
        print("   ✓ 处理器接收到完整的任务数据结构")
        print("\n🚀 客户端使用方式:")
        print("   app.send_task('sandai_operator.process_task',")
        print("                 kwargs={'artifacts': [...], 'options': {...}})")
        print("\n📦 处理器接收格式:")
        print("   {'task_id': '<celery_id>', 'artifacts': [...], 'options': {...}}")
        
    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    return 0


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