import asyncio
import unittest
from collections import defaultdict

from sdax import AsyncTask, AsyncTaskProcessor, TaskContext, TaskFunction

ATTEMPTS = defaultdict(int)


class TestSdaxRetriesAndTimeouts(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        """Reset attempts counter before each test."""
        global ATTEMPTS
        ATTEMPTS = defaultdict(int)

    async def test_retry_logic_succeeds(self):
        """Verify that a task succeeds after a number of retries."""
        processor = AsyncTaskProcessor()
        ctx = TaskContext()

        async def fail_then_succeed(context: TaskContext):
            global ATTEMPTS
            ATTEMPTS["retry_task"] += 1
            if ATTEMPTS["retry_task"] < 3:
                raise ConnectionError("Simulated temporary failure")
            # Success on the 3rd attempt

        task = AsyncTask(
            name="RetryTask",
            execute=TaskFunction(
                function=fail_then_succeed,
                retries=3,  # Allow enough retries to succeed
                backoff_factor=0.1,  # Keep test fast
            ),
        )
        processor.add_task(task, 1)

        await processor.process_tasks(ctx)
        self.assertEqual(ATTEMPTS["retry_task"], 3)

    async def test_retry_logic_fails_after_exhaustion(self):
        """Verify that a task fails permanently if it exceeds its retries."""
        processor = AsyncTaskProcessor()
        ctx = TaskContext()

        async def always_fail(context: TaskContext):
            global ATTEMPTS
            ATTEMPTS["fail_task"] += 1
            raise ConnectionError("Simulated permanent failure")

        task = AsyncTask(
            name="FailTask",
            execute=TaskFunction(function=always_fail, retries=2, backoff_factor=0.1),
        )
        processor.add_task(task, 1)

        with self.assertRaises(ExceptionGroup):
            await processor.process_tasks(ctx)

        # It should try once, then retry twice, for a total of 3 attempts
        self.assertEqual(ATTEMPTS["fail_task"], 3)

    async def test_timeout_is_enforced(self):
        """Verify that a task that takes too long is correctly timed out."""
        processor = AsyncTaskProcessor()
        ctx = TaskContext()

        async def slow_task(context: TaskContext):
            await asyncio.sleep(1)  # This will take too long

        task = AsyncTask(
            name="SlowTask",
            pre_execute=TaskFunction(
                function=slow_task,
                timeout=0.1,  # Set a very short timeout
            ),
        )
        processor.add_task(task, 1)

        with self.assertRaises(ExceptionGroup) as cm:
            await processor.process_tasks(ctx)

        # Check that the underlying exception is indeed a TimeoutError
        self.assertIsInstance(cm.exception.exceptions[0], asyncio.TimeoutError)


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