"""Unit tests for async reward functions.

Tests the async reward function implementation including:
- AsyncORM class functionality
- Event loop daemon thread management
- Performance comparison between sync and async reward functions
"""
import asyncio
import time
import unittest
from typing import List


class TestAsyncRewardFunctions(unittest.TestCase):
    """Test async reward function utilities and base class."""

    def test_start_and_shutdown_event_loop_in_daemon(self):
        """Test that we can start and shutdown an event loop in a daemon thread."""
        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon

        # Start the event loop
        thread, loop, ready_event = start_event_loop_in_daemon(name='TestLoop')

        # Wait for the loop to be ready
        self.assertTrue(ready_event.wait(timeout=5))

        # Verify the loop is running
        self.assertTrue(loop.is_running())
        self.assertTrue(thread.is_alive())

        # Shutdown the event loop
        shutdown_event_loop_in_daemon(thread, loop)

        # Give the thread time to finish
        thread.join(timeout=2)
        self.assertFalse(thread.is_alive())

    def test_run_async_function_in_daemon_loop(self):
        """Test running an async function in the daemon event loop."""
        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon

        thread, loop, ready_event = start_event_loop_in_daemon(name='TestLoop')
        ready_event.wait(timeout=5)

        async def async_add(a, b):
            await asyncio.sleep(0.01)  # Simulate async work
            return a + b

        # Run the async function in the daemon loop
        future = asyncio.run_coroutine_threadsafe(async_add(2, 3), loop)
        result = future.result(timeout=5)

        self.assertEqual(result, 5)

        shutdown_event_loop_in_daemon(thread, loop)

    def test_async_orm_base_class(self):
        """Test that AsyncORM can be subclassed and used correctly."""
        from swift.rewards import AsyncORM

        class TestAsyncReward(AsyncORM):

            async def __call__(self, completions: List[str], **kwargs) -> List[float]:
                await asyncio.sleep(0.01)
                return [float(len(c)) for c in completions]

        reward_func = TestAsyncReward()

        # Check that it's detected as async
        self.assertTrue(asyncio.iscoroutinefunction(reward_func.__call__))

        # Run in an event loop to verify it works
        async def run_test():
            result = await reward_func(['hello', 'world!'])
            return result

        result = asyncio.get_event_loop().run_until_complete(run_test())
        self.assertEqual(result, [5.0, 6.0])

    def test_async_reward_is_detected(self):
        """Test that async reward functions are correctly detected."""
        from swift.rewards import AsyncORM

        class SyncReward:

            def __call__(self, completions, **kwargs):
                return [1.0] * len(completions)

        class AsyncReward(AsyncORM):

            async def __call__(self, completions, **kwargs):
                return [1.0] * len(completions)

        sync_func = SyncReward()
        async_func = AsyncReward()

        # Check detection
        self.assertFalse(asyncio.iscoroutinefunction(sync_func))
        self.assertFalse(asyncio.iscoroutinefunction(sync_func.__call__))
        self.assertTrue(asyncio.iscoroutinefunction(async_func.__call__))


class TestAsyncRewardPerformance(unittest.TestCase):
    """Test performance benefits of async reward functions."""

    def test_parallel_async_execution(self):
        """Test that multiple async reward functions execute in parallel."""
        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon

        thread, loop, ready_event = start_event_loop_in_daemon(name='PerfTestLoop')
        ready_event.wait(timeout=5)

        sleep_time = 0.1  # 100ms per call
        num_calls = 5

        async def slow_async_func(idx):
            await asyncio.sleep(sleep_time)
            return idx

        # Test sequential execution time
        start_seq = time.time()
        for i in range(num_calls):
            future = asyncio.run_coroutine_threadsafe(slow_async_func(i), loop)
            future.result(timeout=5)
        sequential_time = time.time() - start_seq

        # Test parallel execution time using asyncio.gather
        async def run_parallel():
            tasks = [slow_async_func(i) for i in range(num_calls)]
            return await asyncio.gather(*tasks)

        start_par = time.time()
        future = asyncio.run_coroutine_threadsafe(run_parallel(), loop)
        results = future.result(timeout=5)
        parallel_time = time.time() - start_par

        shutdown_event_loop_in_daemon(thread, loop)

        # Verify results
        self.assertEqual(results, list(range(num_calls)))

        # Parallel should be significantly faster than sequential
        # Sequential: ~num_calls * sleep_time = 0.5s
        # Parallel: ~sleep_time = 0.1s
        # Allow some margin for overhead
        self.assertLess(parallel_time, sequential_time * 0.5)

        print('\nPerformance test results:')
        print(f'  Sequential time: {sequential_time:.3f}s (expected ~{sleep_time * num_calls:.1f}s)')
        print(f'  Parallel time:   {parallel_time:.3f}s (expected ~{sleep_time:.1f}s)')
        print(f'  Speedup:         {sequential_time / parallel_time:.1f}x')

    def test_async_reward_function_batch_performance(self):
        """Test performance of async reward function with batch processing."""
        from swift.rewards import AsyncORM
        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon

        sleep_per_item = 0.05  # 50ms per item
        batch_size = 8

        class SlowSyncReward:

            def __call__(self, completions, **kwargs):
                rewards = []
                for c in completions:
                    time.sleep(sleep_per_item)  # Blocking sleep
                    rewards.append(float(len(c)))
                return rewards

        class FastAsyncReward(AsyncORM):

            async def __call__(self, completions, **kwargs):

                async def score_single(text):
                    await asyncio.sleep(sleep_per_item)  # Non-blocking sleep
                    return float(len(text))

                # Process all in parallel
                tasks = [score_single(c) for c in completions]
                return await asyncio.gather(*tasks)

        completions = [f'text_{i}' for i in range(batch_size)]

        # Test sync reward function
        sync_reward = SlowSyncReward()
        start_sync = time.time()
        sync_results = sync_reward(completions)
        sync_time = time.time() - start_sync

        # Test async reward function
        thread, loop, ready_event = start_event_loop_in_daemon(name='BatchPerfLoop')
        ready_event.wait(timeout=5)

        async_reward = FastAsyncReward()

        async def run_async():
            return await async_reward(completions)

        start_async = time.time()
        future = asyncio.run_coroutine_threadsafe(run_async(), loop)
        async_results = future.result(timeout=10)
        async_time = time.time() - start_async

        shutdown_event_loop_in_daemon(thread, loop)

        # Verify results are the same
        self.assertEqual(len(sync_results), len(async_results))
        self.assertEqual(sync_results, list(async_results))

        # Async should be significantly faster
        # Sync: ~batch_size * sleep_per_item = 0.4s
        # Async: ~sleep_per_item = 0.05s
        self.assertLess(async_time, sync_time * 0.5)

        print('\nBatch processing performance:')
        print(f'  Sync time:  {sync_time:.3f}s (expected ~{sleep_per_item * batch_size:.2f}s)')
        print(f'  Async time: {async_time:.3f}s (expected ~{sleep_per_item:.2f}s)')
        print(f'  Speedup:    {sync_time / async_time:.1f}x')


if __name__ == '__main__':
    unittest.main()
