"""Testing the stability of the API in this module"""
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from pxr import Sdf


# ======================================================================
class TestVariantNodes(ogts.OmniGraphTestCase):
    def __init__(self, arg):
        super().__init__(arg)
        self.usda_file = "variant_sets.usda"
        self.input_prim = "/InputPrim"

    @staticmethod
    def set_variant_selection(prim_path, variant_set_name, variant_name):  # noqa: C901
        stage = omni.usd.get_context().get_stage()
        prim = stage.GetPrimAtPath(prim_path)
        variant_set = prim.GetVariantSet(variant_set_name)
        variant_names = variant_set.GetVariantNames()
        if variant_name not in variant_names:
            variant_set.AddVariant(variant_name)
        variant_set.SetVariantSelection(variant_name)

    def set_up_variant_sets(self):
        for prim_path in [self.input_prim]:
            stage = omni.usd.get_context().get_stage()
            prim = stage.DefinePrim(prim_path, "Xform")
            attr = prim.CreateAttribute("primvars:displayColor", Sdf.ValueTypeNames.Color3f)

            variant_set_name = "color"
            variant_sets = prim.GetVariantSets()
            variant_set = variant_sets.GetVariantSet(variant_set_name)

            variant_set.AddVariant("red")
            variant_set.AddVariant("green")
            variant_set.AddVariant("blue")

            variant_set.SetVariantSelection("red")
            with variant_set.GetVariantEditContext():
                attr.Set((1, 0, 0))

            variant_set.SetVariantSelection("green")
            with variant_set.GetVariantEditContext():
                attr.Set((0, 1, 0))

            variant_set.SetVariantSelection("blue")
            with variant_set.GetVariantEditContext():
                attr.Set((0, 0, 1))

            variant_set.SetVariantSelection("default_color")

            variant_set_name = "size"
            variant_sets = prim.GetVariantSets()
            variant_set = variant_sets.GetVariantSet(variant_set_name)

            variant_set.AddVariant("small")
            variant_set.AddVariant("medium")
            variant_set.AddVariant("large")

            variant_set.SetVariantSelection("default_size")

    def get_variant_selection(self, prim_path, variant_set_name):
        stage = omni.usd.get_context().get_stage()
        prim = stage.GetPrimAtPath(prim_path)
        variant_set = prim.GetVariantSet(variant_set_name)
        return variant_set.GetVariantSelection()

    async def _test_variant_node(
        self,
        node_type,
        node_attrs,
        expected_values,
        expected_variants=None,
        expected_attr_values=None,
        read_file=True,
        use_connected_prim=False,
        variant_selection=None,
    ):

        await omni.usd.get_context().new_stage_async()
        if read_file:
            (result, error) = await ogts.load_test_file(self.usda_file, use_caller_subdirectory=True)
            self.assertTrue(result, error)
        else:
            self.set_up_variant_sets()

        if variant_selection:
            prim_path, variant_set_name, variant_name = variant_selection
            self.set_variant_selection(prim_path, variant_set_name, variant_name)

        keys = og.Controller.Keys
        controller = og.Controller()

        graph = "/TestGraph"
        variant_node = "VariantNode"

        values = [(f"{variant_node}.{key}", value) for key, value in node_attrs.items()]

        controller.edit(
            graph,
            {
                keys.CREATE_NODES: [
                    (variant_node, node_type),
                ],
                keys.SET_VALUES: values,
            },
        )

        stage = omni.usd.get_context().get_stage()
        # use a prim path connection instead of setting the value directly
        if use_connected_prim:
            (_, (_prim_at_path, _), _, _) = controller.edit(
                graph,
                {
                    keys.CREATE_NODES: [
                        ("GetPrimAtPath", "omni.graph.nodes.GetPrimsAtPath"),
                        ("ConstToken", "omni.graph.nodes.ConstantToken"),
                    ],
                    keys.CONNECT: [
                        ("ConstToken.inputs:value", "GetPrimAtPath.inputs:path"),
                        ("GetPrimAtPath.outputs:prims", f"{variant_node}.inputs:prim"),
                    ],
                    keys.SET_VALUES: [
                        ("ConstToken.inputs:value", self.input_prim),
                    ],
                },
            )
        else:
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{graph}/{variant_node}.inputs:prim"),
                target=self.input_prim,
            )

        await controller.evaluate()

        for output in expected_values.keys():
            test_attribute = controller.attribute(output, f"{graph}/{variant_node}")
            actual_value = controller.get(attribute=test_attribute)

            if isinstance(expected_values[output], list):
                actual_value = list(actual_value)

            self.assertTrue(
                actual_value == expected_values[output],
                f"actual_value: {actual_value}, expected_value: {expected_values[output]}",
            )

        if expected_variants:
            for prim_path, variant_set_name, expected_variant_name in expected_variants:
                actual_value = self.get_variant_selection(prim_path, variant_set_name)
                self.assertTrue(
                    actual_value == expected_variant_name,
                    f"actual_value: {actual_value}, expected_value: {expected_variant_name}",
                )

        if expected_attr_values:
            for prim_path, attr_name, expected_value in expected_attr_values:
                test_attribute = stage.GetPropertyAtPath(f"{prim_path}.{attr_name}")
                actual_value = test_attribute.Get()
                self.assertTrue(
                    actual_value == expected_value,
                    f"actual_value: {actual_value}, expected_value: {expected_value}",
                )

    async def test_get_variant_names_1(self):
        node_type = "omni.graph.nodes.GetVariantNames"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {"outputs:variantNames": ["blue", "green", "red"]}

        await self._test_variant_node(node_type, node_attrs, expected_values)

    async def test_get_variant_names_2(self):
        node_type = "omni.graph.nodes.GetVariantNames"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {"outputs:variantNames": ["blue", "green", "red"]}

        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=False)
        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=True)

    async def test_get_variant_selection(self):
        node_type = "omni.graph.nodes.GetVariantSelection"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {"outputs:variantName": "default_color"}

        await self._test_variant_node(node_type, node_attrs, expected_values)

    async def test_get_variant_selection_2(self):
        node_type = "omni.graph.nodes.GetVariantSelection"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {"outputs:variantName": "green"}

        variant_selection = (self.input_prim, "color", "green")

        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            variant_selection=variant_selection,
            use_connected_prim=False,
        )
        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            variant_selection=variant_selection,
            use_connected_prim=True,
        )

    async def test_get_variant_set_names(self):
        node_type = "omni.graph.nodes.GetVariantSetNames"
        node_attrs = {}
        expected_values = {"outputs:variantSetNames": ["color", "size"]}

        await self._test_variant_node(node_type, node_attrs, expected_values)

    async def test_get_variant_set_names_2(self):
        node_type = "omni.graph.nodes.GetVariantSetNames"
        node_attrs = {}
        expected_values = {"outputs:variantSetNames": ["color", "size"]}

        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=False)
        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=True)

    async def test_has_variant_set_1(self):
        node_type = "omni.graph.nodes.HasVariantSet"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {"outputs:exists": True}

        await self._test_variant_node(node_type, node_attrs, expected_values)

    async def test_has_variant_set_2(self):
        node_type = "omni.graph.nodes.HasVariantSet"
        node_attrs = {
            "inputs:variantSetName": "flavor",
        }
        expected_values = {"outputs:exists": False}

        await self._test_variant_node(node_type, node_attrs, expected_values)

    async def test_has_variant_set_3(self):
        node_type = "omni.graph.nodes.HasVariantSet"
        node_attrs = {
            "inputs:variantSetName": "flavor",
        }
        expected_values = {"outputs:exists": False}

        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=False)
        await self._test_variant_node(node_type, node_attrs, expected_values, use_connected_prim=True)

    async def test_clear_variant_set_with_default(self):
        node_type = "omni.graph.nodes.ClearVariantSelection"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {}
        expected_variants = [(self.input_prim, "color", "default_color")]

        await self._test_variant_node(node_type, node_attrs, expected_values, expected_variants)

    async def test_clear_variant_set(self):
        node_type = "omni.graph.nodes.ClearVariantSelection"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {}
        expected_variants = [(self.input_prim, "color", "default_color")]

        await self._test_variant_node(node_type, node_attrs, expected_values, expected_variants, read_file=False)

    async def test_clear_variant_set_2(self):
        node_type = "omni.graph.nodes.ClearVariantSelection"
        node_attrs = {
            "inputs:variantSetName": "color",
        }
        expected_values = {}
        expected_variants = [(self.input_prim, "color", "default_color")]

        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            expected_variants,
            read_file=False,
            use_connected_prim=False,
        )
        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            expected_variants,
            read_file=False,
            use_connected_prim=True,
        )

    async def test_set_variant_selection(self):
        node_type = "omni.graph.nodes.SetVariantSelection"
        node_attrs = {
            "inputs:setVariant": True,
            "inputs:variantSetName": "color",
            "inputs:variantName": "red",
        }
        expected_values = {}
        expected_variants = [(self.input_prim, "color", "red")]

        await self._test_variant_node(node_type, node_attrs, expected_values, expected_variants)

    async def test_set_variant_selection_2(self):
        node_type = "omni.graph.nodes.SetVariantSelection"
        node_attrs = {
            "inputs:setVariant": True,
            "inputs:variantSetName": "color",
            "inputs:variantName": "red",
        }
        expected_values = {}
        expected_variants = [(self.input_prim, "color", "red")]

        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            expected_variants,
            use_connected_prim=False,
        )
        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            expected_variants,
            use_connected_prim=True,
        )

    async def test_blend_variants(self):
        node_type = "omni.graph.nodes.BlendVariants"
        node_attrs = {
            "inputs:variantSetName": "color",
            "inputs:blend": 0.5,
            "inputs:variantNameA": "red",
            "inputs:variantNameB": "green",
        }
        expected_values = {}
        expected_variants = []
        expected_attr_values = [(self.input_prim, "primvars:displayColor", [0.5, 0.5, 0.0])]

        await self._test_variant_node(
            node_type,
            node_attrs,
            expected_values,
            expected_variants,
            expected_attr_values,
            read_file=False,
        )
