import omni.kit.test
from .test_base import KitLaunchTest


class TestKitRunLoop(KitLaunchTest):
    async def test_kit_runloop(self):

        # Build kit file
        kit_file = "[dependencies]\n"
        kit_file += '"omni.kit.loop-default" = {}\n'

        kit_file += "[settings.app.runLoops]\n"

        LOOPS = [("loop60", 60), ("loop30", 30), ("main", 300)]
        for loop, freq in LOOPS:
            kit_file += f"{loop}.rateLimitEnabled = true\n"
            kit_file += f"{loop}.rateLimitFrequency = {freq}\n"
            kit_file += f"{loop}.rateLimitUseBusyLoop = true\n"

        # Subscribe to update events, count them. Validate count and thread ids. Measure time. Overall test code should
        # take 1 second = 300 main updates.
        script = f"""
from collections import defaultdict
import threading
import time
import omni.kit.app

app = omni.kit.app.get_app()

def simple_assert(v, message):
    if not v:
        print(f"[error] simple_assert failure: {{message}}")
        omni.kit.app.get_app().post_quit(123)

subs = []
cnt = defaultdict(int)
threads = {{}}
start_t = time.time()
for loop, freq in {LOOPS}:
    def on_update(e, loop=loop):
        cnt[loop] += 1

        # Check thread id
        tid = threading.get_ident()
        if loop not in threads:
            threads[loop] = tid
        simple_assert(threads[loop] == tid, f"runloop: {{loop}}, thread id is incorrect")

        if loop == "main" and cnt[loop] >= 300:
            # validate count
            for loop, freq in {LOOPS}:
                f = cnt[loop] / freq
                simple_assert(f >= 0.80 and f <= 1.20, f"runloop: {{loop}}, expected: {{freq}}, result: {{cnt[loop]}}")

            dt = time.time() - start_t
            simple_assert(dt >= 0.8 and dt <= 1.2, f"test took longer than expected: {{dt}}")
            omni.kit.app.get_app().post_quit(0)

    subs.append(app.get_update_event_stream(loop).create_subscription_to_pop(
        on_update
    ))
        """

        await self._run_kit_with_script(script, args=["--/app/quitAfter=10000"], kit_file=kit_file)
