import omni.kit.test
import omni.graph.core as og
import omni.graph.core.tests as ogts
import os


class TestOgn(ogts.OmniGraphTestCase):

    async def test_generated(self):
        test_data = [{'inputs': [['inputs:name', '_translate', False]], 'outputs': [['outputs:value', {'type': 'pointd[3][]', 'value': [[1, 2, 3], [4, 5, 6]]}, False]], 'setup': {'create_nodes': [['TestNode', 'omni.graph.nodes.GetGatheredAttribute'], ['GatherByPath', 'omni.graph.nodes.GatherByPath']], 'create_prims': [['Empty', {}], ['Xform1', {'_translate': ['pointd[3]', [1, 2, 3]]}], ['Xform2', {'_translate': ['pointd[3]', [4, 5, 6]]}], ['XformTagged1', {'foo': ['token', ''], '_translate': ['pointd[3]', [1, 2, 3]]}], ['Tagged1', {'foo': ['token', '']}]], 'connect': [['GatherByPath.outputs:gatherId', 'TestNode.inputs:gatherId']], 'set_values': [['GatherByPath.inputs:primPaths', ['/Xform1', '/Xform2']], ['GatherByPath.inputs:attributes', '_translate'], ['GatherByPath.inputs:allAttributes', False]]}}]
        test_node = None
        test_graph = None
        for i, test_run in enumerate(test_data):
            inputs = test_run.get('inputs', [])
            outputs = test_run.get('outputs', [])
            state_set = test_run.get('state_set', [])
            state_get = test_run.get('state_get', [])
            setup = test_run.get('setup', None)
            if setup is None or setup:
                await omni.usd.get_context().new_stage_async()
                test_graph = None
            elif not setup:
                self.assertTrue(test_graph is not None and test_graph.is_valid(), "Test is misconfigured - empty setup cannot be in the first test")
            if setup:
                (test_graph, test_nodes, _, _) = og.Controller.edit("/TestGraph", setup)
                self.assertTrue(test_nodes)
                test_node = test_nodes[0]
            elif setup is None:
                if test_graph is None:
                    test_graph = og.Controller.create_graph("/TestGraph")
                    self.assertTrue(test_graph is not None and test_graph.is_valid())
                test_node = og.Controller.create_node(
                    ("TestNode_omni_graph_nodes_GetGatheredAttribute", test_graph), "omni.graph.nodes.GetGatheredAttribute"
                )
            self.assertTrue(test_graph is not None and test_graph.is_valid(), "Test graph invalid")
            self.assertTrue(test_node is not None and test_node.is_valid(), "Test node invalid")
            await og.Controller.evaluate(test_graph)
            values_to_set = inputs + state_set
            if values_to_set:
                for attribute_name, attribute_value, _ in inputs + state_set:
                    og.Controller((attribute_name, test_node)).set(attribute_value)
            await og.Controller.evaluate(test_graph)
            for attribute_name, expected_value, _ in outputs + state_get:
                attribute = og.Controller.attribute(attribute_name, test_node)
                actual_output = og.Controller.get(attribute)
                expected_type = None
                if isinstance(expected_value, dict):
                    expected_type = expected_value["type"]
                    expected_value = expected_value["value"]
                ogts.verify_values(expected_value, actual_output, f"omni.graph.nodes.GetGatheredAttribute User test case #{i+1}: {attribute_name} attribute value error")
                if expected_type:
                    tp = og.AttributeType.type_from_ogn_type_name(expected_type)
                    actual_type = attribute.get_resolved_type()
                    if tp != actual_type:
                        raise ValueError(f"omni.graph.nodes.GetGatheredAttribute User tests - {attribute_name}: Expected {expected_type}, saw {actual_type.get_ogn_type_name()}")

    async def test_data_access(self):
        from omni.graph.nodes.ogn.OgnGetGatheredAttributeDatabase import OgnGetGatheredAttributeDatabase
        test_file_name = "OgnGetGatheredAttributeTemplate.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_GetGatheredAttribute")
        database = OgnGetGatheredAttributeDatabase(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:gatherId"))
        attribute = test_node.get_attribute("inputs:gatherId")
        db_value = database.inputs.gatherId
        expected_value = 0
        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("inputs:name"))
        attribute = test_node.get_attribute("inputs:name")
        db_value = database.inputs.name
        expected_value = ""
        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))
