import os
import omni.kit.test
import omni.graph.core as og
import omni.graph.core.tests as ogts
from omni.graph.core.tests.omnigraph_test_utils import _TestGraphAndNode
from omni.graph.core.tests.omnigraph_test_utils import _test_clear_scene
from omni.graph.core.tests.omnigraph_test_utils import _test_setup_scene
from omni.graph.core.tests.omnigraph_test_utils import _test_verify_scene


class TestOgn(ogts.OmniGraphTestCase):

    TEST_DATA = [
        {
            'inputs': [
                ['inputs:a', {'type': 'float', 'value': 4.1}, False],
            ],
            'outputs': [
                ['outputs:result', 4, False],
            ],
        },
        {
            'inputs': [
                ['inputs:a', {'type': 'half', 'value': -4.9}, False],
            ],
            'outputs': [
                ['outputs:result', -5, False],
            ],
        },
        {
            'inputs': [
                ['inputs:a', {'type': 'double[3]', 'value': [1.3, 2.4, -3.7]}, False],
            ],
            'outputs': [
                ['outputs:result', [1, 2, -4], False],
            ],
        },
        {
            'inputs': [
                ['inputs:a', {'type': 'double[]', 'value': [1.3, 2.4, -3.7, 4.5]}, False],
            ],
            'outputs': [
                ['outputs:result', [1, 2, -4, 4], False],
            ],
        },
    ]

    async def test_generated(self):
        test_info = _TestGraphAndNode()
        controller = og.Controller()
        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            test_info = await _test_setup_scene(self, controller, "/TestGraph", "TestNode_omni_graph_nodes_Floor", "omni.graph.nodes.Floor", test_run, test_info)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.Floor User test case #{i+1}")

    async def test_vectorized_generated(self):
        test_info = _TestGraphAndNode()
        controller = og.Controller()
        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            test_info = await _test_setup_scene(self, controller, "/TestGraph", "TestNode_omni_graph_nodes_Floor","omni.graph.nodes.Floor", test_run, test_info, 16)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.Floor User test case #{i+1}", 16)

    async def test_thread_safety(self):
        import omni.kit
        # Generate multiple instances of the test setup to run them concurrently
        instance_setup = dict()
        for n in range(24):
            instance_setup[f"/TestGraph_{n}"] = _TestGraphAndNode()

        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            for (key, test_info) in instance_setup.items():
                test_info = await _test_setup_scene(self, og.Controller(allow_exists_prim=True), key, "TestNode_omni_graph_nodes_Floor", "omni.graph.nodes.Floor", test_run, test_info)
            self.assertEqual(len(og.get_all_graphs()), 24)

            # We want to evaluate all graphs concurrently. Kick them all.
            # Evaluate multiple times to skip 2 serial frames and increase chances for a race condition.
            for _ in range(10):
                await omni.kit.app.get_app().next_update_async()

            for (key, test_instance) in instance_setup.items():
                _test_verify_scene(self, og.Controller(), test_run, test_info, f"omni.graph.nodes.Floor User test case #{i+1}, instance{key}")

async def test_data_access(self):
    from omni.graph.nodes.ogn.OgnFloorDatabase import OgnFloorDatabase
    test_file_name = "OgnFloorTemplate.usda"
    usd_path = os.path.join(os.path.dirname(__file__), "usd", test_file_name)
    if not os.path.exists(usd_path):
        self.assertTrue(False, f"{usd_path} not found for loading test")
    (result, error) = await ogts.load_test_file(usd_path)
    self.assertTrue(result, f'{error} on {usd_path}')
    test_node = og.Controller.node("/TestGraph/Template_omni_graph_nodes_Floor")
    database = OgnFloorDatabase(test_node)
    self.assertTrue(test_node.is_valid())
    node_type_name = test_node.get_type_name()
    self.assertEqual(og.GraphRegistry().get_node_type_version(node_type_name), 1)

    def _attr_error(attribute: og.Attribute, usd_test: bool) -> str:
        test_type = "USD Load" if usd_test else "Database Access"
        return f"{node_type_name} {test_type} Test - {attribute.get_name()} value error"

