"""Basic tests of the action graph"""
import omni.graph.core as og
import omni.graph.core.tests as ogts


# ======================================================================
class TestActionGraphIntegration(ogts.OmniGraphTestCase):
    """Tests for Script Node"""

    # ----------------------------------------------------------------------
    async def test_compute_creates_dynamic_attrib_legacy(self):
        """Test that the script node can create dynamic attribs within its compute"""
        controller = og.Controller()
        keys = og.Controller.Keys

        script = """
attribute_exists = db.node.get_attribute_exists("inputs:multiplier")
if attribute_exists != True:
    db.node.create_attribute("inputs:multiplier", og.Type(og.BaseDataType.DOUBLE))
db.outputs.data = db.inputs.data * db.inputs.multiplier"""

        (graph, (on_impulse_node, script_node), _, _,) = controller.edit(
            {"graph_path": "/TestGraph", "evaluator_name": "execution"},
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Script", "omni.graph.scriptnode.ScriptNode"),
                ],
                keys.CONNECT: ("OnImpulse.outputs:execOut", "Script.inputs:execIn"),
                keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("Script.inputs:script", script),
                ],
            },
        )

        script_node.create_attribute(
            "inputs:data", og.Type(og.BaseDataType.DOUBLE), og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT
        )
        script_node.create_attribute(
            "outputs:data", og.Type(og.BaseDataType.DOUBLE), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
        )
        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: [
                    ("Script.inputs:data", 42),
                    ("Script.outputs:data", 0),
                ],
            },
        )
        await controller.evaluate(graph)

        # trigger graph evaluation once so the compute runs
        og.Controller.set(controller.attribute("state:enableImpulse", on_impulse_node), True)
        await controller.evaluate(graph)
        val = og.Controller.get(controller.attribute("outputs:data", script_node))
        self.assertEqual(val, 0)

        # set value on the dynamic attrib and check compute
        og.Controller.set(controller.attribute("inputs:multiplier", script_node), 2.0)
        og.Controller.set(controller.attribute("state:enableImpulse", on_impulse_node), True)
        await controller.evaluate(graph)
        val = og.Controller.get(controller.attribute("outputs:data", script_node))
        self.assertEqual(val, 84)

    # ----------------------------------------------------------------------
    async def test_use_loaded_dynamic_attrib(self):
        """Test that the script node can use a dynamic attrib loaded from USD"""
        await ogts.load_test_file("TestScriptNode.usda", use_caller_subdirectory=True)

        controller = og.Controller()
        val = og.Controller.get(controller.attribute("outputs:data", "/World/ActionGraph/script_node"))
        self.assertEqual(val, 0)
        # trigger graph evaluation once so the compute runs
        og.Controller.set(controller.attribute("state:enableImpulse", "/World/ActionGraph/on_impulse_event"), True)
        await controller.evaluate()
        val = og.Controller.get(controller.attribute("outputs:data", "/World/ActionGraph/script_node"))
        self.assertEqual(val, 84)
        as_int = og.Controller.get(controller.attribute("outputs:asInt", "/World/ActionGraph/script_node"))
        as_float = og.Controller.get(controller.attribute("state:asFloat", "/World/ActionGraph/script_node"))
        self.assertEqual(as_int, 84)
        self.assertEqual(as_float, 84.0)

    # ----------------------------------------------------------------------
    async def test_simple_scripts_legacy(self):
        """Test that some simple scripts work as intended"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (script_node,), _, _,) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: ("Script", "omni.graph.scriptnode.ScriptNode"),
            },
        )

        script_node.create_attribute(
            "inputs:my_input_attribute", og.Type(og.BaseDataType.INT), og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT
        )
        script_node.create_attribute(
            "outputs:my_output_attribute", og.Type(og.BaseDataType.INT), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
        )
        output_controller = og.Controller(og.Controller.attribute("outputs:my_output_attribute", script_node))

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", "db.outputs.my_output_attribute = 123"),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 123)

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: [
                    ("Script.inputs:script", "db.outputs.my_output_attribute = -db.inputs.my_input_attribute"),
                    ("Script.inputs:my_input_attribute", 1234),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), -1234)

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: [
                    (
                        "Script.inputs:script",
                        "db.outputs.my_output_attribute = db.inputs.my_input_attribute * db.inputs.my_input_attribute",
                    ),
                    ("Script.inputs:my_input_attribute", -12),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 144)

    # ----------------------------------------------------------------------
    async def test_internal_state_keeps_persistent_info_legacy(self):
        """Test that the script node can keep persistent information using internal state"""

        script = """
if (not hasattr(db.internal_state, 'num1')):
    db.internal_state.num1 = 0
    db.internal_state.num2 = 1
else:
    sum = db.internal_state.num1 + db.internal_state.num2
    db.internal_state.num1 = db.internal_state.num2
    db.internal_state.num2 = sum
db.outputs.data = db.internal_state.num1"""

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

        (graph, (_, script_node), _, _,) = controller.edit(
            {"graph_path": "/TestGraph", "evaluator_name": "execution"},
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Script", "omni.graph.scriptnode.ScriptNode"),
                ],
                keys.CONNECT: ("OnTick.outputs:tick", "Script.inputs:execIn"),
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Script.inputs:script", script),
                ],
            },
        )

        script_node.create_attribute(
            "outputs:data", og.Type(og.BaseDataType.INT), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
        )
        output_controller = og.Controller(og.Controller.attribute("outputs:data", script_node))

        # Check that the script node produces the Fibonacci numbers
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 0)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 1)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 1)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 2)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 3)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 5)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 8)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 13)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 21)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 34)
        await controller.evaluate(graph)
        self.assertEqual(output_controller.get(), 55)

    # ----------------------------------------------------------------------
    async def test_script_global_scope(self):
        """Test that variables, functions, and classes defined outside of the user-defined callbacks are visible"""

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

        (graph, (script_node,), _, _,) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: ("Script", "omni.graph.scriptnode.ScriptNode"),
            },
        )

        script_node.create_attribute(
            "outputs:output_a", og.Type(og.BaseDataType.INT), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
        )
        script_node.create_attribute(
            "outputs:output_b", og.Type(og.BaseDataType.TOKEN), og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
        )
        output_a_controller = og.Controller(og.Controller.attribute("outputs:output_a", script_node))
        output_b_controller = og.Controller(og.Controller.attribute("outputs:output_b", script_node))

        script_test_constants = """
MY_CONSTANT_A = 123
MY_CONSTANT_B = 'foo'

def setup(db):
    db.outputs.output_a = MY_CONSTANT_A

def compute(db):
    db.outputs.output_b = MY_CONSTANT_B"""

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", script_test_constants),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_a_controller.get(), 123)
        self.assertEqual(output_b_controller.get(), "foo")
        await controller.evaluate(graph)
        self.assertEqual(output_b_controller.get(), "foo")

        script_test_variables = """
my_variable = 123

def setup(db):
    global my_variable
    my_variable = 234
    db.outputs.output_a = my_variable

def compute(db):
    global my_variable
    db.outputs.output_b = f'{my_variable}'
    my_variable += 1"""

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", script_test_variables),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_a_controller.get(), 234)
        self.assertEqual(output_b_controller.get(), "234")
        await controller.evaluate(graph)
        self.assertEqual(output_b_controller.get(), "235")

        script_test_functions = """
def my_function_a():
    return 123

my_variable_b = 'foo'

def my_function_b():
    return my_variable_b

def setup(db):
    db.outputs.output_a = my_function_a()

def compute(db):
    db.outputs.output_b = my_function_b()
    global my_variable_b
    my_variable_b = 'bar'"""

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", script_test_functions),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_a_controller.get(), 123)
        self.assertEqual(output_b_controller.get(), "foo")
        await controller.evaluate(graph)
        self.assertEqual(output_b_controller.get(), "bar")

        script_test_imports = """
import inspect
import math

my_lambda = lambda x: x
code_len = len(inspect.getsource(my_lambda))

def setup(db):
    db.outputs.output_a = code_len

def compute(db):
    db.outputs.output_b = f'{math.pi:.2f}'"""

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", script_test_imports),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_a_controller.get(), 24)
        self.assertEqual(output_b_controller.get(), "3.14")
        await controller.evaluate(graph)
        self.assertEqual(output_b_controller.get(), "3.14")

        script_test_classes = """
class MyClass:
    def __init__(self, value):
        self.value = value

    def get_value(self):
        return self.value

    @staticmethod
    def get_num():
        return 123

my_variable = MyClass('foo')

def setup(db):
    db.outputs.output_a = MyClass.get_num()

def compute(db):
    db.outputs.output_b = my_variable.get_value()
    my_variable.value = 'bar'"""

        controller.edit(
            "/TestGraph",
            {
                keys.SET_VALUES: ("Script.inputs:script", script_test_classes),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(output_a_controller.get(), 123)
        self.assertEqual(output_b_controller.get(), "foo")
        await controller.evaluate(graph)
        self.assertEqual(output_b_controller.get(), "bar")
