import utils
import asyncio
import threading


class BasicUsage(utils.AergiaUnitTestCase):

    def test_asyncless(self):
        def a():
            x = 100
            while x > 0:
                x -= 1

        self.Aergia.start()
        a()
        self.Aergia.stop()

        samples = self.Aergia.get_samples()
        self.assertFalse(samples)

    def test_sequential_tasks(self):
        delay = 0.2
        num_times = 5

        async def b(tot, num):
            await asyncio.sleep(delay)
            return tot + num

        async def a():
            tot = 0
            for i in range(num_times):
                tot = await b(tot, i)
            assert tot == 10

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()

        self.assertFuncContains('b',
                                [self.expected_samples(delay * num_times)],
                                samples)

    def test_simultaneous_tasks(self):
        delay = 0.2
        async def b(): await asyncio.sleep(delay)
        async def a(): await asyncio.gather(b(), b(), b())

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()

        self.assertFuncContains('b', [self.expected_samples(delay * 3)],
                                samples)
        # TODO samples from gather all execution time, should we trace this?
        self.assertFuncContains('a', [self.expected_samples(delay)], samples)

    def test_subthread_task(self):
        delay = 0.2

        async def b(): await asyncio.sleep(delay)
        async def a(): await asyncio.gather(b(), b(), b())
        def c(): asyncio.run(a())

        self.Aergia.start()
        x = threading.Thread(target=c)
        x.start()
        x.join()
        self.Aergia.stop()

        samples = self.Aergia.get_samples()
        self.assertFuncContains('c', [], samples)
        self.assertFuncContains('b', [self.expected_samples(delay * 3)],
                                samples)
        # TODO samples from gather all execution time, should we trace this?
        self.assertFuncContains('a', [self.expected_samples(delay)], samples)

    def test_eager_task(self):
        delay = 0.2

        async def a():
            proc = await asyncio.create_subprocess_shell(f'sleep {delay}')
            await proc.communicate()

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()
        self.assertFuncContains('a', [self.expected_samples(delay)], samples)

    def test_async_generator(self):
        delay = 0.2
        num_times = 10

        async def b():
            for i in range(num_times):
                await asyncio.sleep(delay)
                yield i

        async def a():
            lst = []
            async for item in b():
                lst.append(item)

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()
        # TODO can we make these results more intuitive?
        # async generators also do not work at all with yappi.
        # doing so would be unique to this profiler.

        # self.assertFuncContains('b',
        #                         [self.expected_samples(delay * num_times)],
        #                         samples)
        # self.assertFuncContains('a', [], samples)

        self.assertFuncContains('a',
                                [self.expected_samples(delay * num_times)],
                                samples)
        self.assertFuncContains('b', [], samples)

    def test_async_gen_and_comp(self):
        delay = 0.2
        num_times = 10

        async def b():
            for i in range(num_times):
                await asyncio.sleep(delay)
                yield i

        async def a():
            return [r async for r in b()]

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()
        # TODO can we make these results more intuitive?
        # async generators also do not work at all with yappi.
        # doing so would be unique to this profiler.

        self.assertFuncContains('b', [], samples)
        self.assertFuncContains('a', [], samples)
        self.assertFuncContains('<listcomp>',
                                [self.expected_samples(delay * num_times)],
                                samples)

    def test_deep_await(self):
        delay = 0.2

        async def c(): await asyncio.sleep(delay)
        async def b(): await c()
        async def a(): await b()

        self.Aergia.start()
        asyncio.run(a())
        self.Aergia.stop()

        samples = self.Aergia.get_samples()

        self.assertFuncContains('c', [self.expected_samples(delay)], samples)
        self.assertFuncContains('b', [], samples)
        self.assertFuncContains('a', [], samples)