"""Support for generating simple regression test code for OmniGraph Nodes.

The tests do three things:
    1. Do a test "import" of the Python interface to the node
    2. Do a test load of the USDA template interface for the node
    3. Run all of the "tests" cases specified in the node's .ogn file

Exports:
    generate_tests: Create a TestNODE.py file containing standard unit tests of the node operation
"""
from pathlib import Path
from typing import List, Optional

from .attributes.management import list_without_runtime_attributes
from .generate_test_imports import ensure_test_is_imported
from .keys import MemoryTypeValues, TestKeys
from .nodes import NodeInterfaceGenerator
from .utils import GeneratorConfiguration, UnimplementedError, logger

__all__ = ["generate_tests"]


class NodeTestGenerator(NodeInterfaceGenerator):
    """Manage the functions required to generate basic test scripts for a node"""

    def __init__(self, configuration: GeneratorConfiguration):  # noqa: PLW0246
        """Set up the generator and output the test scripts for the node

        Just passes the initialization on to the parent class. See the argument and exception descriptions there.
        """
        super().__init__(configuration)

    # ----------------------------------------------------------------------
    def interface_file_name(self) -> str:
        """Return the path to the name of the Python test file"""
        return self.test_class_name() + ".py"

    # ----------------------------------------------------------------------
    def test_class_name(self) -> str:
        """Returns the name to use for the test class and base name for the test file"""
        return f"Test{self.base_name}"

    # ----------------------------------------------------------------------
    def generate_user_test_data(self, all_tests: List) -> bool:
        """Generate the section of the user test that creates the test data to be iterated over.

        The variable "test_data" is set up to be a list of test data, whose elements consist of four sub-lists
            - values for input attributes, set before the test starts
            - values for output attributes, checked after the test finishes
            - initial values for state attributes, set before the test starts
            - final values for state attributes, checked after the test finishes
        Each of the data elements consist of a 3-tuple of the name of the attribute, the value for it, and whether the
        attribute's data should be manipulated on the GPU.

        Simple example of one test where two inputs are set, one output is checked, one state attribute has an
        initial value set and a final value checked, and a node of type "MyNodeType" is created before the test runs:
            test_data = [
                {
                    "inputs": [
                        ["inputs:attr", INPUT_VALUE, False],
                        ["inputs:attr2", INPUT_VALUE2, False],
                    ],
                    "outputs": [
                        ["outputs:attr", OUTPUT_VALUE, False],
                    ],
                    "state_set": [
                        ["state:attr", INITIAL_VALUE, False],
                    ],
                    "state_get": [
                        ["state:attr", FINAL_VALUE, False],
                    ],
                    "setup": {
                        "create_nodes": [["MyNode", "MyNodeType"]]
                    }
                }
            ]

        Args:
            all_tests: List of test dictionaries with the test information to be consolidated

        Returns:
            True if the node tests use the old V1 setup data, meant for the deprecated OmniGraphHelper
        """
        test_data = []
        uses_v1_setup = False
        for node_test_data in all_tests:
            # The outputs who decide on being on the GPU at runtime have to be marked as such in the test data.
            # That way the test knows where to look for the results.
            gpu_outputs = node_test_data.gpu_outputs
            test_run = {}
            # Do all inputs first, for clarity
            if node_test_data.input_values:
                input_data = []
                for attribute, attribute_value in node_test_data.input_values.items():
                    if attribute.memory_type == MemoryTypeValues.ANY:
                        UnimplementedError(f"Input '{attribute.name}' has unsupported memory CPU/GPU")
                    on_gpu = attribute.memory_type != MemoryTypeValues.CPU
                    input_data.append([attribute.name, attribute_value, on_gpu])
                test_run[TestKeys.INPUTS] = input_data
            # ...then all outputs
            if node_test_data.expected_outputs:
                output_data = []
                for attribute, attribute_value in node_test_data.expected_outputs.items():
                    if attribute.memory_type == MemoryTypeValues.ANY:
                        on_gpu = attribute.name in gpu_outputs
                    else:
                        on_gpu = attribute.memory_type != MemoryTypeValues.CPU
                    output_data.append([attribute.name, attribute_value, on_gpu])
                test_run[TestKeys.OUTPUTS] = output_data
            # ...then all state initial values
            if node_test_data.state_initial_values:
                state_data = []
                for attribute, attribute_value in node_test_data.state_initial_values.items():
                    if attribute.memory_type == MemoryTypeValues.ANY:
                        on_gpu = attribute.name in gpu_outputs
                    else:
                        on_gpu = attribute.memory_type != MemoryTypeValues.CPU
                    state_data.append([attribute.name, attribute_value, on_gpu])
                test_run[TestKeys.STATE_SET] = state_data
            # ... all state final values
            if node_test_data.state_final_values:
                state_data = []
                for attribute, attribute_value in node_test_data.state_final_values.items():
                    if attribute.memory_type == MemoryTypeValues.ANY:
                        on_gpu = attribute.name in gpu_outputs
                    else:
                        on_gpu = attribute.memory_type != MemoryTypeValues.CPU
                    state_data.append([attribute.name, attribute_value, on_gpu])
                test_run[TestKeys.STATE_GET] = state_data
            # ... and graph setup (where None has a different meaning than the empty list)
            if node_test_data.graph_setup is not None:
                test_run[TestKeys.SETUP] = node_test_data.graph_setup
                uses_v1_setup |= node_test_data.uses_v1_setup
            test_data.append(test_run)

        self.out.write(f"test_data = {test_data}")
        return uses_v1_setup

    # ----------------------------------------------------------------------
    def generate_user_test_run(self, uses_v1_setup: bool = False):
        """Generate the section of the user test that iterates test runs over the test data"""
        node_type_name = self.node_interface.name
        # Emit the test loop, creating a new node each iteration to ensure that all unset values are their default
        self.out.write("test_node = None")
        self.out.write("test_graph = None")
        if self.out.indent("for i, test_run in enumerate(test_data):"):
            self.out.write(f"inputs = test_run.get('{TestKeys.INPUTS}', [])")
            self.out.write(f"outputs = test_run.get('{TestKeys.OUTPUTS}', [])")
            self.out.write(f"state_set = test_run.get('{TestKeys.STATE_SET}', [])")
            self.out.write(f"state_get = test_run.get('{TestKeys.STATE_GET}', [])")
            self.out.write(f"setup = test_run.get('{TestKeys.SETUP}', None)")

            # Clean out the stage for the next test, unless the setup is to be carried on
            if self.out.indent("if setup is None or setup:"):
                self.out.write("await omni.usd.get_context().new_stage_async()")
                self.out.write("test_graph = None")
                self.out.exdent()
            if self.out.indent("elif not setup:"):
                # An empty setup means "use the previous setup", which should have created a graph and a node
                self.out.write(
                    "self.assertTrue(test_graph is not None and test_graph.is_valid(), "
                    + '"Test is misconfigured - empty setup cannot be in the first test")'
                )
                self.out.exdent()

            # Emit the construction of the test node
            if self.out.indent("if setup:"):
                if uses_v1_setup:
                    self.out.write("await ogts.setup_test_environment()")
                    self.out.write('test_nodes = og.OmniGraphHelper().edit_graph("/", setup)')
                    self.out.write("test_graph = og.get_current_graph()")
                else:
                    self.out.write('(test_graph, test_nodes, _, _) = og.Controller.edit("/TestGraph", setup)')
                self.out.write("self.assertTrue(test_nodes)")
                self.out.write("test_node = test_nodes[0]")
                self.out.exdent()
            if self.out.indent("elif setup is None:"):
                # A previous setup will have created a test_graph, otherwise create a default graph for testing.
                if self.out.indent("if test_graph is None:"):
                    self.out.write('test_graph = og.Controller.create_graph("/TestGraph")')
                    self.out.write("self.assertTrue(test_graph is not None and test_graph.is_valid())")
                    self.out.exdent()
                if self.out.indent("test_node = og.Controller.create_node("):
                    self.out.write(f'("TestNode_{self.safe_name()}", test_graph), "{node_type_name}"')
                    self.out.exdent(")")
                self.out.exdent()
            self.out.write('self.assertTrue(test_graph is not None and test_graph.is_valid(), "Test graph invalid")')
            self.out.write('self.assertTrue(test_node is not None and test_node.is_valid(), "Test node invalid")')

            # First evaluation sets up the node and Fabric
            self.out.write("await og.Controller.evaluate(test_graph)")

            # Emit the code that sets the inputs of the node first
            self.out.write("values_to_set = inputs + state_set")
            if self.out.indent("if values_to_set:"):
                if self.out.indent("for attribute_name, attribute_value, _ in inputs + state_set:"):
                    self.out.write("og.Controller((attribute_name, test_node)).set(attribute_value)")
                    self.out.exdent()
                self.out.exdent()

            # Emit code to evaluate the node
            self.out.write("await og.Controller.evaluate(test_graph)")

            # Emit code to read the outputs and state and compare against the expected values
            if self.out.indent("for attribute_name, expected_value, _ in outputs + state_get:"):
                self.out.write("attribute = og.Controller.attribute(attribute_name, test_node)")
                self.out.write("actual_output = og.Controller.get(attribute)")
                self.out.write("expected_type = None")
                if self.out.indent("if isinstance(expected_value, dict):"):
                    self.out.write('expected_type = expected_value["type"]')
                    self.out.write('expected_value = expected_value["value"]')
                    self.out.exdent()
                error_message = f"{node_type_name} User test case #{{i+1}}: {{attribute_name}} attribute value error"
                self.out.write(f'ogts.verify_values(expected_value, actual_output, f"{error_message}")')
                if self.out.indent("if expected_type:"):
                    self.out.write("tp = og.AttributeType.type_from_ogn_type_name(expected_type)")
                    self.out.write("actual_type = attribute.get_resolved_type()")
                    if self.out.indent("if tp != actual_type:"):
                        self.out.write(
                            f'raise ValueError(f"{node_type_name} User tests - {{attribute_name}}: '
                            f'Expected {{expected_type}}, saw {{actual_type.get_ogn_type_name()}}")'
                        )
                        self.out.exdent()
                    self.out.exdent()
                self.out.exdent()

            self.out.exdent()
        self.out.exdent()

    # ----------------------------------------------------------------------
    def generate_user_tests(self):
        """Generate the test method that exercises the tests specified in the .ogn file"""
        all_tests = self.node_interface.all_tests()
        # If the user did not specify any tests then do not generate this test function
        if not all_tests:
            return

        self.out.write("")
        if self.out.indent("async def test_generated(self):"):
            uses_v1_setup = self.generate_user_test_data(all_tests)
            self.generate_user_test_run(uses_v1_setup)

    # ----------------------------------------------------------------------
    def generate_data_access_test(self):
        """Generate the test method for loading the generated USD file"""
        check_usd = "usd" not in self.node_interface.excluded_generators
        check_python = "python" not in self.node_interface.excluded_generators
        # There is not enough information to test data access if there's no Python database and no USD file
        if not check_usd and not check_python:
            return

        node_type_name = self.node_interface.name
        node_name = f"Template_{self.safe_name()}"

        # The node may have generated a .usda file with default values on it. If so then load it in and confirm
        # that the node exists after load and the inputs have the defaults.
        self.out.write("")
        self.out.write("async def test_data_access(self):")
        if self.out.indent():
            if check_python:
                db_name = f"{self.base_name}Database"
                self.out.write(f"from {self.module}.ogn.{db_name} import {db_name}")
            # If USD testing is turned on then use that as a source of the test node
            if check_usd:
                self.out.write(f'test_file_name = "{self.base_name}Template.usda"')
                # The file is found in the usd/ subdirectory of this script's test directory.
                self.out.write('usd_path = os.path.join(os.path.dirname(__file__), "usd", test_file_name)')
                if self.out.indent("if not os.path.exists(usd_path):"):
                    self.out.write('self.assertTrue(False, f"{usd_path} not found for loading test")')
                    self.out.exdent()
                self.out.write("(result, error) = await ogts.load_test_file(usd_path)")
                self.out.write("self.assertTrue(result, f'{error} on {usd_path}')")
                # This node name is hardcoded into the USD generation
                self.out.write(f'test_node = og.Controller.node("/TestGraph/{node_name}")')
            # otherwise create the node directly
            else:
                self.out.write('(_, (test_node,), _, _) = og.Controller.edit("/TestGraph", {')
                self.out.write(f'    og.Controller.Keys.CREATE_NODES: ("{node_name}", "{node_type_name}")')
                self.out.write("})")

            if check_python:
                self.out.write(f"database = {db_name}(test_node)")
            self.out.write("self.assertTrue(test_node.is_valid())")

            # Checking the version instead of the name allows for node type name aliases
            self.out.write("node_type_name = test_node.get_type_name()")
            version = self.node_interface.version
            self.out.write(f"self.assertEqual(og.GraphRegistry().get_node_type_version(node_type_name), {version})")
            self.out.write("")
            if self.out.indent("def _attr_error(attribute: og.Attribute, usd_test: bool) -> str:"):
                self.out.write('test_type = "USD Load" if usd_test else "Database Access"')
                self.out.write('return f"{node_type_name} {test_type} Test - {attribute.get_name()} value error"')
                self.out.exdent()
            self.out.write("")

            # Emit code to read the inputs and compare against the default values
            for attribute in list_without_runtime_attributes(self.node_interface.all_attributes()):
                # Optional attributes are not written to the template file
                if not attribute.is_required:
                    continue
                name_to_check = attribute.usd_name()
                self.out.write("")
                self.out.write(f'self.assertTrue(test_node.get_attribute_exists("{name_to_check}"))')

                # Always get the value, so that the code is exercised
                self.out.write(f'attribute = test_node.get_attribute("{name_to_check}")')
                # Values aren't available for testing if the memory type is purely CUDA
                if check_python and attribute.memory_type != MemoryTypeValues.CUDA:
                    attribute_accessor = f"{attribute.namespace}.{attribute.python_property_name()}"
                    # If memory type is determined at runtime the property is an accessor, not a value
                    if attribute.memory_type == MemoryTypeValues.ANY and not attribute.cpp_accessor_on_cpu():
                        attribute_accessor += ".cpu"
                    self.out.write(f"db_value = database.{attribute_accessor}")

                # Only check input numbers since they should be well-defined as the defaults for all nodes
                if not attribute.is_read_only():
                    continue

                expected_value = attribute.python_default_value()
                if expected_value is not None:
                    self.out.write(f"expected_value = {expected_value}")
                    if check_usd:
                        self.out.write("actual_value = og.Controller.get(attribute)")
                        self.out.write("ogts.verify_values(expected_value, actual_value, _attr_error(attribute, True))")
                    if check_python and attribute.memory_type != MemoryTypeValues.CUDA:
                        self.out.write("ogts.verify_values(expected_value, db_value, _attr_error(attribute, False))")
            self.out.exdent()

    # ----------------------------------------------------------------------
    def generate_node_interface(self):
        """Generate the test method for the named node"""
        self.generate_user_tests()
        self.generate_data_access_test()
        if self.interface_directory is not None:
            ensure_test_is_imported(self.test_class_name(), Path(self.interface_directory))

    # ----------------------------------------------------------------------
    def pre_interface_generation(self):
        """Create the imports and common test framework used by tests for all nodes"""
        self.out.write("import omni.kit.test")
        self.out.write("import omni.graph.core as og")
        self.out.write("import omni.graph.core.tests as ogts")
        if "usd" not in self.node_interface.excluded_generators:
            self.out.write("import os")
        self.out.write("")
        self.out.write("")
        self.out.write("class TestOgn(ogts.OmniGraphTestCase):")
        self.out.indent()


# ======================================================================
def generate_tests(configuration: GeneratorConfiguration) -> Optional[str]:
    """Create support files for the tests on the node

    Args:
        configuration: Information defining how and where the test files will be generated

    Returns:
        String containing the generated test script code or None if its generation was not enabled

    Raises:
        NodeGenerationError: When there is a failure in the generation of the tests
    """
    if not configuration.node_interface.can_generate("tests"):
        return None

    logger.info("Generating Node Type Tests")
    generator = NodeTestGenerator(configuration)
    generator.generate_interface()
    return str(generator.out)
