#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
并发获取密集库位处理器测试
验证优化后的处理器在高并发场景下的安全性
"""

import asyncio
import pytest
import uuid
from typing import Dict, Any, List
from unittest.mock import Mock, patch, AsyncMock
from datetime import datetime

# 假设的测试环境导入
from services.execution.handlers.storage_location import GetIdleCrowdedSiteBlockHandler
from services.execution.task_context import TaskContext
from data.models.operate_point_layer import OperatePointLayer
from data.models.storage_area import StorageArea


class TestConcurrentStorageLocation:
    """并发获取密集库位测试类"""
    
    @pytest.fixture
    def mock_session(self):
        """模拟数据库会话"""
        session = AsyncMock()
        session.execute = AsyncMock()
        session.commit = AsyncMock()
        session.rollback = AsyncMock()
        return session
    
    @pytest.fixture
    def mock_context(self):
        """模拟任务上下文"""
        context = Mock(spec=TaskContext)
        context.map_id = "test_map_001"
        context.task_record_id = "task_001"
        context.block_record_id = "block_001"
        context.set_variable = Mock()
        context.set_block_output = Mock()
        context.is_task_canceled_async = AsyncMock(return_value=False)
        return context
    
    @pytest.fixture
    def handler(self):
        """获取密集库位处理器实例"""
        return GetIdleCrowdedSiteBlockHandler()
    
    @pytest.fixture
    def mock_storage_area(self):
        """模拟库区数据"""
        area = Mock(spec=StorageArea)
        area.area_name = "test_area"
        area.select_logic = 0  # 默认按库位名称排序
        return area
    
    @pytest.fixture
    def mock_storage_layers(self):
        """模拟库位数据"""
        layers = []
        for i in range(1, 6):  # 创建5个库位
            layer = Mock(spec=OperatePointLayer)
            layer.layer_name = f"A00{i}"
            layer.area_name = "test_area"
            layer.scene_id = "test_map_001"
            layer.is_disabled = False
            layer.is_deleted = False
            layer.is_locked = False
            layer.is_occupied = False  # 默认空闲
            layer.goods_content = ""
            layer.goods_stored_at = None
            layer.last_access_at = datetime.now()
            layers.append(layer)
        return layers
    
    async def simulate_concurrent_requests(self, handler, num_requests: int, input_params: Dict[str, Any], mock_context):
        """模拟并发请求"""
        tasks = []
        
        for i in range(num_requests):
            # 每个请求使用不同的任务ID
            context = Mock(spec=TaskContext)
            context.map_id = mock_context.map_id
            context.task_record_id = f"task_{i:03d}"
            context.block_record_id = f"block_{i:03d}"
            context.set_variable = Mock()
            context.set_block_output = Mock()
            context.is_task_canceled_async = AsyncMock(return_value=False)
            
            block = {"name": f"test_block_{i}"}
            
            # 创建异步任务
            task = asyncio.create_task(
                handler.execute(block, input_params.copy(), context)
            )
            tasks.append(task)
        
        # 等待所有任务完成
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return results
    
    @patch('services.execution.handlers.storage_location.get_async_session')
    async def test_concurrent_lock_requests(self, mock_get_session, handler, mock_context, 
                                          mock_storage_area, mock_storage_layers):
        """测试并发锁定请求 - 应该只有一个成功"""
        
        # 设置模拟数据库会话
        mock_session = AsyncMock()
        mock_get_session.return_value.__aenter__.return_value = mock_session
        
        # 模拟查询库区信息
        area_result = AsyncMock()
        area_result.scalar_one_or_none.return_value = mock_storage_area
        
        # 模拟查询锁定库位（无锁定库位）
        locked_result = AsyncMock()
        locked_result.fetchall.return_value = []
        
        # 模拟查询所有库位
        all_layers_result = AsyncMock()
        all_layers_result.scalars.return_value.all.return_value = mock_storage_layers
        
        # 模拟查询候选库位
        candidate_result = AsyncMock()
        candidate_result.scalars.return_value.all.return_value = mock_storage_layers
        
        # 设置查询结果的顺序
        mock_session.execute.side_effect = [
            area_result,          # 查询库区
            locked_result,        # 查询锁定库位
            all_layers_result,    # 查询所有库位
            candidate_result,     # 查询候选库位
        ]
        
        # 模拟原子更新操作 - 第一个请求成功，后续失败
        update_results = []
        successful_updates = 1  # 只允许一个成功
        
        def mock_execute_update(*args, **kwargs):
            nonlocal successful_updates
            result = AsyncMock()
            if successful_updates > 0:
                result.rowcount = 1  # 更新成功
                successful_updates -= 1
            else:
                result.rowcount = 0  # 更新失败（被抢占）
            return result
        
        # 输入参数
        input_params = {
            "groupName": '["test_area"]',
            "filled": "false",  # 放货
            "content": "test_goods",
            "lock": "true",     # 需要锁定
            "retry": "false"
        }
        
        # 启动多个并发请求
        num_requests = 5
        
        # 重置模拟对象状态
        mock_session.reset_mock()
        
        # 为每个并发请求设置独立的模拟
        concurrent_tasks = []
        for i in range(num_requests):
            task_context = Mock(spec=TaskContext)
            task_context.map_id = "test_map_001"
            task_context.task_record_id = f"task_{i:03d}"
            task_context.block_record_id = f"block_{i:03d}"
            task_context.set_variable = Mock()
            task_context.set_block_output = Mock()
            task_context.is_task_canceled_async = AsyncMock(return_value=False)
            
            # 为每个请求创建独立的会话模拟
            task_session = AsyncMock()
            
            # 为每个会话设置查询结果
            task_area_result = AsyncMock()
            task_area_result.scalar_one_or_none.return_value = mock_storage_area
            
            task_locked_result = AsyncMock()
            task_locked_result.fetchall.return_value = []
            
            task_all_layers_result = AsyncMock()
            task_all_layers_result.scalars.return_value.all.return_value = mock_storage_layers
            
            task_candidate_result = AsyncMock()
            task_candidate_result.scalars.return_value.all.return_value = mock_storage_layers
            
            task_session.execute.side_effect = [
                task_area_result,
                task_locked_result,
                task_all_layers_result,
                task_candidate_result,
                mock_execute_update()  # 原子更新操作
            ]
            
            block = {"name": f"test_block_{i}"}
            
            # 创建独立的会话上下文管理器
            async def get_task_session():
                return task_session
            
            # 动态patch每个任务的session
            with patch('services.execution.handlers.storage_location.get_async_session') as mock_task_session:
                mock_task_session.return_value.__aenter__.return_value = task_session
                task = asyncio.create_task(
                    handler.execute(block, input_params.copy(), task_context)
                )
                concurrent_tasks.append(task)
        
        # 等待所有任务完成
        results = await asyncio.gather(*concurrent_tasks, return_exceptions=True)
        
        # 验证结果
        successful_results = [r for r in results if isinstance(r, dict) and r.get("success")]
        failed_results = [r for r in results if isinstance(r, dict) and not r.get("success")]
        
        # 断言：应该只有一个请求成功
        assert len(successful_results) == 1, f"应该只有1个请求成功，实际成功了 {len(successful_results)} 个"
        assert len(failed_results) == num_requests - 1, f"应该有 {num_requests-1} 个请求失败，实际失败了 {len(failed_results)} 个"
        
        # 验证成功的请求返回了正确的库位
        successful_result = successful_results[0]
        assert successful_result["data"]["siteId"] == "A001"
        assert successful_result["data"]["areaName"] == "test_area"
        assert successful_result["data"]["locked"] == True
    
    @patch('services.execution.handlers.storage_location.get_async_session')
    async def test_concurrent_no_lock_requests(self, mock_get_session, handler, mock_context,
                                             mock_storage_area, mock_storage_layers):
        """测试并发非锁定请求 - 可能有多个成功但不会冲突"""
        
        # 设置模拟数据库会话
        mock_session = AsyncMock()
        mock_get_session.return_value.__aenter__.return_value = mock_session
        
        # 为非锁定场景，所有原子更新都应该成功
        def mock_execute_update(*args, **kwargs):
            result = AsyncMock()
            result.rowcount = 1  # 所有更新都成功
            return result
        
        # 输入参数（不锁定）
        input_params = {
            "groupName": '["test_area"]',
            "filled": "false",  # 放货
            "content": "test_goods",
            "lock": "false",    # 不锁定
            "retry": "false"
        }
        
        num_requests = 3
        results = await self.simulate_concurrent_requests(handler, num_requests, input_params, mock_context)
        
        # 在不锁定的情况下，由于我们的原子更新机制，仍然应该保证并发安全
        # 每个请求仍然需要通过原子检查占用状态
        successful_results = [r for r in results if isinstance(r, dict) and r.get("success")]
        
        # 所有请求都应该能够成功（因为我们模拟了足够的空闲库位）
        assert len(successful_results) >= 1, "至少应该有一个请求成功"
    
    async def test_performance_benchmark(self, handler, mock_context):
        """性能基准测试 - 测试优化后的性能"""
        
        input_params = {
            "groupName": '["test_area"]',
            "filled": "false",
            "content": "test_goods",
            "lock": "true",
            "retry": "false"
        }
        
        # 模拟大量并发请求
        num_requests = 50
        
        start_time = asyncio.get_event_loop().time()
        
        with patch('services.execution.handlers.storage_location.get_async_session'):
            results = await self.simulate_concurrent_requests(handler, num_requests, input_params, mock_context)
        
        end_time = asyncio.get_event_loop().time()
        execution_time = end_time - start_time
        
        print(f"并发处理 {num_requests} 个请求耗时: {execution_time:.3f} 秒")
        print(f"平均每个请求耗时: {execution_time/num_requests*1000:.2f} 毫秒")
        
        # 性能断言（根据实际需求调整）
        assert execution_time < 10.0, f"处理 {num_requests} 个并发请求不应超过10秒，实际耗时 {execution_time:.3f} 秒"
    
    async def test_retry_mechanism_under_concurrency(self, handler, mock_context):
        """测试重试机制在并发环境下的表现"""
        
        input_params = {
            "groupName": '["test_area"]',
            "filled": "false",
            "content": "test_goods",
            "lock": "true",
            "retry": "true",
            "retryNum": "3",
            "retryPeriod": "100"  # 100ms重试间隔
        }
        
        # 模拟中等并发
        num_requests = 10
        
        with patch('services.execution.handlers.storage_location.get_async_session'):
            results = await self.simulate_concurrent_requests(handler, num_requests, input_params, mock_context)
        
        # 验证重试机制工作正常
        successful_results = [r for r in results if isinstance(r, dict) and r.get("success")]
        
        # 在重试机制下，成功率应该有所提升
        success_rate = len(successful_results) / num_requests
        print(f"重试机制下的成功率: {success_rate*100:.1f}%")
        
        # 断言至少有一定的成功率
        assert success_rate >= 0.1, f"重试机制下成功率过低: {success_rate*100:.1f}%"


# 运行测试的主函数
async def run_concurrent_tests():
    """运行并发测试"""
    print("=== 开始并发获取密集库位处理器测试 ===")
    
    test_instance = TestConcurrentStorageLocation()
    
    # 创建必要的fixtures
    handler = test_instance.handler()
    mock_context = test_instance.mock_context()
    mock_storage_area = test_instance.mock_storage_area()
    mock_storage_layers = test_instance.mock_storage_layers()
    
    try:
        print("\n1. 测试并发锁定请求...")
        await test_instance.test_concurrent_lock_requests(
            None, handler, mock_context, mock_storage_area, mock_storage_layers
        )
        print("✓ 并发锁定请求测试通过")
        
        print("\n2. 测试并发非锁定请求...")
        await test_instance.test_concurrent_no_lock_requests(
            None, handler, mock_context, mock_storage_area, mock_storage_layers
        )
        print("✓ 并发非锁定请求测试通过")
        
        print("\n3. 性能基准测试...")
        await test_instance.test_performance_benchmark(handler, mock_context)
        print("✓ 性能基准测试通过")
        
        print("\n4. 重试机制并发测试...")
        await test_instance.test_retry_mechanism_under_concurrency(handler, mock_context)
        print("✓ 重试机制并发测试通过")
        
        print("\n=== 所有测试完成 ===")
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        raise


if __name__ == "__main__":
    # 直接运行测试
    asyncio.run(run_concurrent_tests())