import omni.kit.test
import gzip
import json
import os

from omni.kit.core.tests import KitLaunchTest

PYTHON_EVENT_PREFIX = "Py::"

class TestProfilePython(KitLaunchTest):
    async def test_profile_python(self):

        FUNC_1 = "some_rare_magical_function_27026930"
        FUNC_2 = "another_rare_magical_function_120941920"

        # Call some functions
        script = f"""
def {FUNC_2}():
    pass

def {FUNC_1}():
    {FUNC_2}()

{FUNC_1}()
        """

        args = [
            # "--/app/quitAfter=10"
            "--enable",
            "omni.kit.profile_python",
        ]

        trace_file = f"{self._temp_folder}/trace.gz"

        args += [
            "--/plugins/carb.profiler-cpu.plugin/saveProfile=1",
            "--/plugins/carb.profiler-cpu.plugin/compressProfile=1",
            "--/app/profileFromStart=1",
            f"--/plugins/carb.profiler-cpu.plugin/filePath='{trace_file}'",
        ]

        # run kit with python profiler and this script
        await self._run_kit_with_script(script, args=args)

        self.assertTrue(os.path.exists(trace_file))

        # Parse trace and check that they have those functions at least mentioned there
        events = set()
        with gzip.open(trace_file, 'rb') as f_in:
            j = json.loads(f_in.read().decode("utf-8") )
            for entry in j:
                name = entry["name"]
                if name.startswith(PYTHON_EVENT_PREFIX):
                    events.add(name)


        self.assertTrue(PYTHON_EVENT_PREFIX + FUNC_1 in events)
        self.assertTrue(PYTHON_EVENT_PREFIX + FUNC_2 in events)
