"""
Contains support for testing methods in ../autonode/autonode.py
"""
import enum
import random
import sys
import unittest
from typing import Tuple

import omni.graph.core as og
import omni.graph.core.autonode as oga
import omni.graph.core.tests as ogts

# Test warp integration


# Detect the testing environment and use the Kit Async testing framework if it is available
try:
    import omni.kit.test

    TestBaseClass = omni.kit.test.AsyncTestCase
except ImportError:
    TestBaseClass = unittest.TestCase


@oga.AutoClass(module_name="__test__")
class TestClass:
    """Simple AutoClass test with four attributes"""

    test_visible_str_property: str = "test visible str property"
    test_visible_int_property: int = 42
    test_visible_float_property: float = 3.141
    test_uninitialized_property: str

    _test_hidden_state = 0

    def __init__(self):
        raise RuntimeError("This is a test singleton class")

    @classmethod
    def _get_instance(cls):
        if not hasattr(cls, "_instance") or cls._instance is None:
            cls._instance = cls.__new__(cls)
        return cls._instance

    def test_visible_method(self, is_input: bool = False) -> int:
        self._test_hidden_state ^= int(is_input)
        return self._test_hidden_state


@oga.AutoFunc(module_name="__test__")
def get_test_class_instance() -> TestClass:
    return TestClass._get_instance()  # noqa: PLW0212


class TestEnum(enum.Enum):
    VAL0 = 0
    VAL1 = 1
    VAL2 = 2


def get_test_enum(input_value: int) -> TestEnum:
    i = input_value % 3
    return TestEnum(i)


@oga.AutoFunc(module_name="__test__", pure=True)
def pack(v: og.Float3, c: og.Color3f) -> og.Bundle:
    bundle = og.Bundle("return", False)
    bundle.create_attribute("vec", og.Float3).value = v
    bundle.create_attribute("col", og.Color3f).value = c
    return bundle


@oga.AutoFunc(module_name="__test__", pure=True)
def unpack_vector(bundle: og.Bundle) -> og.Float3:
    vec = bundle.attribute_by_name("vec")
    if vec:
        return vec.value
    return (0, 0, 0)


@oga.AutoFunc(module_name="__test__", pure=True)
def identity_float3(input_value: og.Float3) -> og.Float3:
    return input_value


@oga.AutoFunc(module_name="__test__", pure=True)
def identity_tuple(input_value_1: str, input_value_2: int) -> Tuple[int, str]:
    return (input_value_1, input_value_2)


def dbl_to_matd(input_value: og.Double3) -> og.Matrix4d:
    pass


################################################################################


class TestAutographNodes(ogts.OmniGraphTestCase):

    TEST_GRAPH_PATH = "/World/TestGraph"

    async def setUp(self):
        """Set up  test environment, to be torn down when done"""
        await super().setUp()
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

    # --------------------------------------------------------------------------------
    async def _property_test(self, property_name, test_value, *, almost_equal=False):
        """exercise a basic network of execution nodes"""

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

        gci_name = "GetClassInstance"
        getter_name = "Getter"
        setter_name = "Setter"

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Ticker", "omni.graph.action.OnTick"),
                    (f"{gci_name}01", "__test__.get_test_class_instance"),
                    (getter_name, f"__test__.TestClass__PROP__{property_name}__GET"),
                ],
                keys.CONNECT: [
                    ("Ticker.outputs:tick", f"{gci_name}01.inputs:exec"),
                    (f"{gci_name}01.outputs:exec", f"{getter_name}.inputs:exec"),
                    (f"{gci_name}01.outputs:out_0", f"{getter_name}.inputs:target"),
                ],
                keys.SET_VALUES: [("Ticker.inputs:onlyPlayback", False)],
            },
        )
        await controller.evaluate(graph)

        (_, _, getter_node) = nodes

        if almost_equal:
            self.assertAlmostEqual(
                og.Controller.get(controller.attribute("outputs:out_0", getter_node)),
                getattr(TestClass._get_instance(), property_name),  # noqa: PLW0212
                places=3,
            )
        else:
            self.assertEqual(
                og.Controller.get(controller.attribute("outputs:out_0", getter_node)),
                getattr(TestClass._get_instance(), property_name),  # noqa: PLW0212
            )

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:exec", getter_node)), og.ExecutionAttributeState.ENABLED
        )

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    (f"{gci_name}02", "__test__.get_test_class_instance"),
                    (setter_name, f"__test__.TestClass__PROP__{property_name}__SET"),
                ],
                keys.CONNECT: [
                    ("Ticker.outputs:tick", f"{gci_name}02.inputs:exec"),
                    (f"{gci_name}02.outputs:exec", f"{setter_name}.inputs:exec"),
                    (f"{gci_name}02.outputs:out_0", f"{setter_name}.inputs:target"),
                ],
                keys.DISCONNECT: ("Ticker.outputs:tick", f"{gci_name}01.inputs:exec"),
                keys.SET_VALUES: (f"{setter_name}.inputs:value", test_value),
            },
        )
        await controller.evaluate(graph)

        (_, setter_node) = nodes

        if almost_equal:
            self.assertAlmostEqual(
                getattr(TestClass._get_instance(), property_name), test_value, places=3  # noqa: PLW0212
            )
        else:
            self.assertEqual(getattr(TestClass._get_instance(), property_name), test_value)  # noqa: PLW0212
        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:exec", setter_node)), og.ExecutionAttributeState.ENABLED
        )

    # --------------------------------------------------------------------------------
    async def test_property_string(self):
        """Test to see strings can be get and set"""
        property_name = "test_visible_str_property"
        test_value = "test_value!"
        await self._property_test(property_name, test_value)

    # --------------------------------------------------------------------------------
    async def test_property_int(self):
        """Test to see ints can be get and set"""
        property_name = "test_visible_int_property"
        test_value = random.randint(0, sys.maxsize >> 32)
        await self._property_test(property_name, test_value)

    # --------------------------------------------------------------------------------
    async def test_property_float(self):
        """Test to see floats can be get and set"""
        property_name = "test_visible_float_property"
        test_value = random.random()
        await self._property_test(property_name, test_value, almost_equal=True)

    # --------------------------------------------------------------------------------
    async def test_enum(self):
        """Test switching on enums"""

        oga.AutoClass(module_name="__test__")(TestEnum)
        oga.AutoFunc(module_name="__test__")(get_test_enum)

        test_enum_sanitized_name = TestEnum.__qualname__.replace(".", "__NSP__")

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

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Ticker", "omni.graph.action.OnTick"),
                    ("Generator", "__test__.get_test_enum"),
                    ("Switch", f"__test__.{test_enum_sanitized_name}"),
                ],
                keys.CONNECT: [
                    ("Ticker.outputs:tick", "Generator.inputs:exec"),
                    ("Generator.outputs:exec", "Switch.inputs:exec"),
                    ("Generator.outputs:out_0", "Switch.inputs:enum"),
                ],
                keys.SET_VALUES: [("Ticker.inputs:onlyPlayback", False), ("Generator.inputs:input_value", 1)],
            },
        )
        await controller.evaluate(graph)

        (_, _, switch_node) = nodes

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL0", switch_node)), og.ExecutionAttributeState.DISABLED
        )

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL1", switch_node)), og.ExecutionAttributeState.ENABLED
        )

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL2", switch_node)), og.ExecutionAttributeState.DISABLED
        )

        controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.SET_VALUES: [("Generator.inputs:input_value", 2), ("Switch.outputs:VAL1", False)]},
        )
        await controller.evaluate(graph)

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL0", switch_node)), og.ExecutionAttributeState.DISABLED
        )

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL1", switch_node)), og.ExecutionAttributeState.DISABLED
        )

        self.assertEqual(
            og.Controller.get(controller.attribute("outputs:VAL2", switch_node)), og.ExecutionAttributeState.ENABLED
        )

    # --------------------------------------------------------------------------------


class TestAutographPureNodes(ogts.OmniGraphTestCase):

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    async def tuple_test(self):
        """Test multiple outputs"""

        def switch_io(input1: str, input2: int) -> Tuple[int, str]:
            return (input2, input1)

        oga.AutoFunc(module_name="__test__", pure=True)(switch_io)

    # ----------------------------------------------------------------------
    async def test_tuple(self):
        """Test multiple outputs"""

        test_input1 = "test"
        test_input2 = 42

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

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [("Identity1", "__test__.identity_tuple"), ("Identity2", "__test__.identity_tuple")],
                keys.CONNECT: [
                    ("Identity1.outputs:out_0", "Identity2.inputs:input_value_1"),
                    ("Identity1.outputs:out_1", "Identity2.inputs:input_value_2"),
                ],
                keys.SET_VALUES: [
                    ("Identity1.inputs:input_value_1", test_input1),
                    ("Identity1.inputs:input_value_2", test_input2),
                ],
            },
        )
        await controller.evaluate(graph)

        (identity1_node, _) = nodes
        self.assertEqual(og.Controller.get(controller.attribute("outputs:out_0", identity1_node)), test_input1)
        self.assertEqual(og.Controller.get(controller.attribute("outputs:out_1", identity1_node)), test_input2)

    # --------------------------------------------------------------------------------
    async def test_bundle(self):
        """Test switching on bundles"""

        test_input_1 = (random.random(), random.random(), random.random())

        test_input_2 = (0, 0.5, 1)

        oga.AutoFunc(module_name="__test__", pure=True)(unpack_vector)

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

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Pack", "__test__.pack"),
                    ("UnpackVector", "__test__.unpack_vector"),
                    ("Id", "__test__.identity_float3"),
                ],
                keys.CONNECT: [
                    # outputs_out_0 instead of outputs:out_0 due to the way bundles are implemented
                    ("Pack.outputs_out_0", "UnpackVector.inputs:bundle"),
                    ("UnpackVector.outputs:out_0", "Id.inputs:input_value"),
                ],
                keys.SET_VALUES: [("Pack.inputs:v", test_input_1), ("Pack.inputs:c", test_input_2)],
            },
        )
        await controller.evaluate(graph)

        (_, unpack_vector_node, _) = nodes
        out_0 = og.Controller.get(controller.attribute("outputs:out_0", unpack_vector_node))
        self.assertAlmostEqual(out_0[0], test_input_1[0], places=3)
        self.assertAlmostEqual(out_0[1], test_input_1[1], places=3)
        self.assertAlmostEqual(out_0[2], test_input_1[2], places=3)

    # --------------------------------------------------------------------------------

    async def test_type(self):
        """Tests creating nodes based on functions from many types."""

        oga.AutoFunc(module_name="__test__")(dbl_to_matd)
