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:array', {'type': 'int[]', 'value': []}, False],
                ['inputs:removeAll', False, False],
                ['inputs:value', {'type': 'int', 'value': 41}, False],
            ],
            'outputs': [
                ['outputs:found', False, False],
                ['outputs:array', {'type': 'int[]', 'value': []}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'int[]', 'value': [41, 42, 41]}, False],
                ['inputs:removeAll', False, False],
                ['inputs:value', {'type': 'int', 'value': 41}, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'int[]', 'value': [42, 41]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'int[]', 'value': [41, 42, 43, 41]}, False],
                ['inputs:removeAll', True, False],
                ['inputs:value', {'type': 'int', 'value': 41}, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'int[]', 'value': [42, 43]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'int[]', 'value': [41, 41, 41, 41]}, False],
                ['inputs:removeAll', True, False],
                ['inputs:value', {'type': 'int', 'value': 41}, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'int[]', 'value': []}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'token[]', 'value': ['FOOD', 'BARFOOD', 'BAZ']}, False],
                ['inputs:value', {'type': 'token', 'value': 'FOO'}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', False, False],
                ['outputs:array', {'type': 'token[]', 'value': ['FOOD', 'BARFOOD', 'BAZ']}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'token[]', 'value': ['FOOD', 'BARFOOD', 'BAZ']}, False],
                ['inputs:value', {'type': 'token', 'value': 'FOOD'}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'token[]', 'value': ['BARFOOD', 'BAZ']}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'int64[]', 'value': [41, 42]}, False],
                ['inputs:value', {'type': 'int64', 'value': 41}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'int64[]', 'value': [42]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'uchar[]', 'value': [41, 42]}, False],
                ['inputs:value', {'type': 'uchar', 'value': 41}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'uchar[]', 'value': [42]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'uint[]', 'value': [41, 42]}, False],
                ['inputs:value', {'type': 'uint', 'value': 42}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'uint[]', 'value': [41]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'uint64[]', 'value': [41, 42]}, False],
                ['inputs:value', {'type': 'uint64', 'value': 42}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'uint64[]', 'value': [41]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'bool[]', 'value': [False, True]}, False],
                ['inputs:value', {'type': 'bool', 'value': False}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'bool[]', 'value': [True]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'half[2][]', 'value': [[1, 2], [3, 4]]}, False],
                ['inputs:value', {'type': 'half[2]', 'value': [1, 2]}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'half[2][]', 'value': [[3, 4]]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'double[2][]', 'value': [[1, 2], [3, 4], [5, 6]]}, False],
                ['inputs:value', {'type': 'double[2]', 'value': [1, 2]}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'double[2][]', 'value': [[3, 4], [5, 6]]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'double[2][]', 'value': [[1, 2], [3, 4], [5, 6], [1, 2]]}, False],
                ['inputs:value', {'type': 'double[2]', 'value': [1, 2]}, False],
                ['inputs:removeAll', True, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'double[2][]', 'value': [[3, 4], [5, 6]]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'double[3][]', 'value': [[1.1, 2.1, 1.0], [3, 4, 5]]}, False],
                ['inputs:value', {'type': 'double[3]', 'value': [3, 4, 5]}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'double[3][]', 'value': [[1.1, 2.1, 1.0]]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'double[3][]', 'value': [[1.1, 2.1, 1.0], [3, 4, 5]]}, False],
                ['inputs:value', {'type': 'double[3]', 'value': [3.1, 4, 5]}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', False, False],
                ['outputs:array', {'type': 'double[3][]', 'value': [[1.1, 2.1, 1.0], [3, 4, 5]]}, False],
            ],
        },
        {
            'inputs': [
                ['inputs:array', {'type': 'int[4][]', 'value': [[1, 2, 3, 4], [3, 4, 5, 6]]}, False],
                ['inputs:value', {'type': 'int[4]', 'value': [1, 2, 3, 4]}, False],
                ['inputs:removeAll', False, False],
            ],
            'outputs': [
                ['outputs:found', True, False],
                ['outputs:array', {'type': 'int[4][]', 'value': [[3, 4, 5, 6]]}, 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_ArrayRemoveValue", "omni.graph.nodes.ArrayRemoveValue", test_run, test_info)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.ArrayRemoveValue 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_ArrayRemoveValue","omni.graph.nodes.ArrayRemoveValue", test_run, test_info, 16)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.ArrayRemoveValue 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_ArrayRemoveValue", "omni.graph.nodes.ArrayRemoveValue", 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.ArrayRemoveValue User test case #{i+1}, instance{key}")

async def test_data_access(self):
    from omni.graph.nodes.ogn.OgnArrayRemoveValueDatabase import OgnArrayRemoveValueDatabase
    test_file_name = "OgnArrayRemoveValueTemplate.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_ArrayRemoveValue")
    database = OgnArrayRemoveValueDatabase(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"


    self.assertTrue(test_node.get_attribute_exists("inputs:removeAll"))
    attribute = test_node.get_attribute("inputs:removeAll")
    db_value = database.inputs.removeAll
    expected_value = False
    actual_value = og.Controller.get(attribute)
    ogts.verify_values(expected_value, actual_value, _attr_error(attribute, True))
    ogts.verify_values(expected_value, db_value, _attr_error(attribute, False))

    self.assertTrue(test_node.get_attribute_exists("outputs:found"))
    attribute = test_node.get_attribute("outputs:found")
    db_value = database.outputs.found
