# Copyright 2024 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import asyncio
import timeit
import unittest

from iree.runtime import (
    get_device,
    HalDeviceLoopBridge,
)


class HalDeviceLoopBridgeTest(unittest.TestCase):
    def testBridge(self):
        loop = asyncio.new_event_loop()
        bridge = HalDeviceLoopBridge(self.device, loop)
        sem1 = None
        sem2 = None
        report = None

        async def main():
            def done_1(x):
                report("PYTHON: sem2.signal(1)")
                sem2.signal(1)

            def done_2(x):
                report("PYTHON: sem2.signal(2)")
                sem2.signal(2)

            f1 = bridge.on_semaphore(sem1, 1, "Semaphore 1 Signaled")
            f1.add_done_callback(done_1)
            f2 = bridge.on_semaphore(sem2, 1, "Semaphore 2 Signaled")
            f2.add_done_callback(done_2)
            f2_again = bridge.on_semaphore(sem2, 2, "Semaphore 2 Signaled Again")

            sem1.signal(1)
            f1_result = await f1
            report("PYTHON: await f1 =", f1_result)
            f2_result = await f2
            report("PYTHON: await f2 =", f2_result)
            f2_again_result = await f2_again
            report("PYTHON: await f2_again =", f2_again_result)

            self.assertEqual(f1_result, "Semaphore 1 Signaled")
            self.assertEqual(f2_result, "Semaphore 2 Signaled")
            self.assertEqual(f2_again_result, "Semaphore 2 Signaled Again")
            report("PYTHON: ASYNC MAIN() COMPLETE")

        def run_iter(with_report):
            nonlocal sem1
            nonlocal sem2
            nonlocal report
            sem1 = self.device.create_semaphore(0)
            sem2 = self.device.create_semaphore(0)
            if with_report:
                report = lambda *args: print(*args)
            else:
                report = lambda *args: None
            loop.run_until_complete(main())

        try:
            run_iter(True)
            iter_time = timeit.timeit("run_iter(False)", globals=locals(), number=10)
            print(f"Time/iter = {iter_time}s")
        finally:
            bridge.stop()

    def setUp(self):
        super().setUp()
        self.device = get_device("local-task")
        self.allocator = self.device.allocator


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