"""Backward compatible module for omni.graph.nodes version 1.15.
Revert your access to the earlier package to this version:

    import omni.graph.nodes._1_15 as on

Or import both and cherry-pick which versions of specific APIs you want to use

    import omni.graph.nodes as on
    import omni.graph.nodes._1_15 as on1_15

    if i_want_v1_15:
        on1_15.old_function()
    else:
        on.old_function()
"""
import ast
import json
from tempfile import NamedTemporaryFile
from typing import Any, Dict, List, Optional, Tuple, Union
from unittest import TestCase

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.graph.tools as ogt
from omni.kit.app import get_app
from pxr import Usd

__all__ = [
    "bundle_inspector_outputs",
    "bundle_inspector_values_as_dictionary",
    "compare_bundle_contents",
    "compare_bundle_inspector_results",
    "computed_bundle_values",
    "create_prim_with_everything",
    "graph_node_path",
    "inspected_bundle_values",
    "inspector_node_definition",
    "prim_node_definition",
    "set_up_test_stage",
]


# ======================================================================
# Type of data coming from a bundle inspector output
BundleInspectorOutputType = Tuple[int, List[str], List[str], List[str], List[int], List[int], List[Any]]


# ======================================================================
# Attribute information for bundle inspector nodes
BUNDLE_INSPECTOR_ATTRIBUTES = [
    ("uint64", "outputs:count"),
    ("token[]", "outputs:names"),
    ("token[]", "outputs:types"),
    ("token[]", "outputs:roles"),
    ("int[]", "outputs:arrayDepths"),
    ("int[]", "outputs:tupleCounts"),
    ("token[]", "outputs:values"),
    ("bundle", "outputs_bundle"),
    ("bool", "inputs:print"),
]


# ======================================================================
# Feed data for constructing sample data with all attribute types in it. This is a list of
# (ATTRIBUTE_TYPE, ATTRIBUTE_SAMPLE_VALUE) that can be used to populate bundles. Tuple counts are
# lifted from the sample value type information and arrays are constructed explicitly from this simple data.
# (Be careful with the half values; they have lower resolution and are better tested in fractionaly powers of 2.)
# fmt: off
PRIM_SAMPLE_VALUES = [
    ("int", 1),
    ("int64", 2),
    ("half", 3.5),
    ("float", 5.6),
    ("double", 7.8),
    ("bool", 1),
    ("token", "goofy"),
    ("uchar", 9),
    ("uint", 10),
    ("uint64", 11),
    ("int", (12, -13)),
    ("half", (14.5, -16.5)),
    ("float", (18.19, 20.21)),
    ("double", (22.23, 24.25)),
    ("int", (26, -27, 28)),
    ("half", (29.5, -31.5, 33.5)),
    ("float", (35.36, -37.38, 39.40)),
    ("double", (41.42, -43.44, 45.46)),
    ("int", (47, -48, 49, -50)),
    ("half", (51.5, -53.5, 55.5, -57.5)),
    ("float", (59.60, -61.62, 63.64, -65.66)),
    ("double", (67.68, -69.70, 71.72, -73.74)),
    ("quath", (75.5, -77.5, 79.5, -81.5)),
    ("quatf", (83.84, -85.86, 87.88, -89.90)),
    ("quatd", (91.92, -93.94, 95.96, -97.98)),
    ("matrixd", ((1.1, -1.2), (2.1, -2.2))),
    ("matrixd", ((1.1, -1.2, 1.3), (2.1, -2.2, 2.3), (3.1, -3.2, 3.3))),
    ("matrixd", ((1.1, -1.2, 1.3, 1.4), (2.1, -2.2, 2.3, -2.4), (3.1, -3.2, 3.3, -3.4), (4.1, -4.2, 4.3, -4.4))),
]
# fmt: on


# ======================================================================
# Data types used for the node description conversions (set as simple type definitions because you can differentiate
# between the two of them just by checking size)
PrimAttributeValueType = List[Union[Tuple[str, Any], Tuple[str, str, Any]]]
OmniGraphNodeDefinitionType = Tuple[str, str, List[str], List[Tuple[str, str]]]
PrimNodeDefinitionType = Tuple[str, Dict[str, Tuple[str, Any, int]]]


# ======================================================================
@ogt.deprecated_function("Use og.Controller.create_prim()")
def prim_node_definition(prim_name: str, data: PrimAttributeValueType, as_usd: bool = True) -> PrimNodeDefinitionType:
    """Returns the data as a USD definition compatible with what is needed for node definitions in set_up_test_stage

    Args:
        prim_name: Name of the prim to be defined
        data: List of prim attribute values. The members of the list come in two forms:
            (TYPE, VALUE) : Creates attributes and arrays of those attributes with names derived from the type
                roughly ("float", 3) -> "float floatAttr = 3.0", "float[] floatAttrArray = [3.0, 3.0, 3.0]"
            (TYPE, NAME, VALUE): Creates an attribute with the given name, type, and value
                roughly ("float", "myFloat", 3) -> "float myFloat = 3.0"
        as_usd: If True then create the definition with USD ordering (i.e. quaternions reordered)
    """
    definition = {}
    for attribute_info in data:
        if len(attribute_info) == 2:
            (attribute_type, value) = attribute_info
            if isinstance(value, str) or (not isinstance(value, list) and not isinstance(value, tuple)):
                tuple_count = 1
                attribute_name = f"{attribute_type}Attr"
            else:
                tuple_count = len(value)
                attribute_name = f"{attribute_type}{tuple_count}Attr"
            # In all of the user-facing data the quaternions are specified as I,J,K,R, however in the USD file and
            # quaternion constructor they expect R,I,J,K so reorder for consistency
            if as_usd and attribute_type.find("quat") == 0:
                the_value = (value[3], value[0], value[1], value[2])
            else:
                the_value = value
            definition[attribute_name] = (f"{attribute_type}", the_value, False, tuple_count)
            definition[f"{attribute_name}Array"] = (f"{attribute_type}", [the_value] * 3, True, tuple_count)
        else:
            (attribute_type, attribute_name, value) = attribute_info
            is_array = False
            the_value = value
            if isinstance(value, str) or (not isinstance(value, list) and not isinstance(value, tuple)):
                tuple_count = 1
            elif isinstance(value, list) and value:
                is_array = True
                if isinstance(value[0], str) or (not isinstance(value[0], list) and not isinstance(value[0], tuple)):
                    tuple_count = 1
                else:
                    tuple_count = len(value[0])
                    # In all of the user-facing data the quaternions are specified as I,J,K,R, however in the USD file
                    # and quaternion constructor they expect R,I,J,K so reorder for consistency
                    if as_usd and attribute_type.find("quat") == 0:
                        the_value = tuple((member[3], member[0], member[1], member[2]) for member in value)
            elif isinstance(value, tuple):
                tuple_count = len(value)
                # In all of the user-facing data the quaternions are specified as I,J,K,R, however in the USD file and
                # quaternion constructor they expect R,I,J,K so reorder for consistency
                if as_usd and attribute_type.find("quat") == 0:
                    the_value = (value[3], value[0], value[1], value[2])
            definition[attribute_name] = (f"{attribute_type}", the_value, is_array, tuple_count)
    return (prim_name, definition)


# ======================================================================
@ogt.deprecated_function("Use og.Controller.create_prim()")
def _write_usd_prim_definition(attribute_information: PrimNodeDefinitionType, out):
    """Write out a USD Prim definition using the attribute information

    Args:
        attribute_information: Tuple(prim_name, Dict of ATTRIBUTE_NAME: (USD_TYPE, USD_VALUE, TUPLE_COUNT))
            e.g. a float[3] named foo defaulted to zeroes would be { "foo": ("float", [0.0, 0.0, 0.0], 3) }
        out: Output object to which the definition will be written
    """
    out.write(f'def Output "{attribute_information[0]}"')
    if out.indent("{"):
        for attribute_name, (type_name, attribute_value, is_array, tuple_count) in attribute_information[1].items():
            usd_type = ogt.usd_type_name(type_name, tuple_count, is_array)
            value = attribute_value
            if usd_type.startswith("token") or usd_type.startswith("string"):
                value = json.dumps(attribute_value)
            out.write(f"{usd_type} {attribute_name} = {value}")
        out.exdent("}")


# ======================================================================
@ogt.deprecated_function("Use og.Controller.create_node()")
def inspector_node_definition(node_name: str, source_bundle: str) -> Tuple[str, str, List[str]]:
    """Returns the entry in a node definition list for a bundle inspector node connected to the given source"""
    return (
        node_name,
        "omni.graph.nodes.BundleInspector",
        [f"rel inputs:bundle = [<{source_bundle}>]"] + ogt.attributes_as_usd(BUNDLE_INSPECTOR_ATTRIBUTES),
        BUNDLE_INSPECTOR_ATTRIBUTES,
    )


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.bundle_inspector_results()")
def inspected_bundle_values(bundle_data) -> List[Any]:
    """Returns a list of attribute values expected on the bundle inspector when applied to the bundle data.
    The list ordering should correspond to what is in BUNDLE_INSPECTOR_ATTRIBUTES, for efficiency.

    Args:
        bundle_data: Raw output of all the output attributes the bundle inspector provides
    """
    names = []
    type_names = []
    role_names = []
    array_flags = []
    tuple_counts = []
    values = []
    # Sorting by name is the easiest way to produce a predictable ordering for the expected bundle values
    for name in sorted(bundle_data.keys()):
        (raw_type_name, value, is_array, tuple_count) = bundle_data[name]
        names.append(name)
        (type_name, role_name) = ogt.separate_ogn_role_and_type(raw_type_name)
        type_names.append(type_name)
        role_names.append(role_name)
        array_flags.append(is_array)
        tuple_counts.append(tuple_count * tuple_count if role_name in ["frame", "matrix", "transform"] else tuple_count)
        values.append(value)
    return [len(bundle_data), names, type_names, role_names, array_flags, tuple_counts, values]


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.bundle_inspector_results()")
def computed_bundle_values(inspector_attributes: List[Any]) -> BundleInspectorOutputType:
    """Returns attribute values computed from a bundle inspector in the same form as inspected_bundle_values

    Args:
        inspector_attributes: List of the Attribute objects on the node in the order of BUNDLE_INSPECTOR_ATTRIBUTES
    """
    # Have to keep these separated because the first attribute is a single value while the rest are lists of strings
    size = og.Controller.get(inspector_attributes[0])
    if size == 0:
        return (0, [], [], [], [], [], [])
    results = []
    for attribute_index in range(1, 7):
        attribute_values = og.Controller.get(inspector_attributes[attribute_index])
        if attribute_index == 6:
            # Bundled values are formatted as strings but are actually a mixture of types
            attribute_values = [
                ast.literal_eval(value) if value != "__unsupported__" else "__unsupported__"
                for value in attribute_values
            ]
        results.append(attribute_values)
    # If these aren't sorted it will be more difficult to compare them
    zipped = list(zip(results[0], results[1], results[2], results[3], results[4], results[5]))
    zipped.sort()
    results = list(zip(*zipped))
    return tuple([size] + results)


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.bundle_inspector_results()")
def bundle_inspector_outputs(bundle_inspector: og.NODE_TYPE_HINTS) -> BundleInspectorOutputType:
    """Returns attribute values computed from a bundle inspector in the same form as inspected_bundle_values

    Args:
        bundle_inspector: Node from which the output is to be extracted
    """
    raw_values = [
        og.Controller.get(og.Controller.attribute(name, bundle_inspector))
        for (attr_type, name) in BUNDLE_INSPECTOR_ATTRIBUTES
        if attr_type != "bundle" and name.startswith("outputs:")
    ]
    size = raw_values[0]
    if size == 0:
        return (0, [], [], [], [], [], [])

    results = []
    for attribute_name in [name for (_, name) in BUNDLE_INSPECTOR_ATTRIBUTES[1:] if name.startswith("outputs:")]:
        results.append(og.Controller.get(og.Controller.attribute(attribute_name, bundle_inspector)))

    # Bundled values are formatted as strings but are actually a mixture of types
    results[-1] = [
        ast.literal_eval(value) if value != "__unsupported__" else "__unsupported__" for value in results[-1]
    ]

    # If these aren't sorted it will be more difficult to compare them
    zipped = list(zip(results[0], results[1], results[2], results[3], results[4], results[5]))
    zipped.sort()
    results = [list(x) for x in zip(*zipped)]
    return tuple([size] + results)


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.bundle_inspector_results()")
def bundle_inspector_values_as_dictionary(inspected_values: List[Any]) -> Dict[str, Tuple[str, str, bool, int, Any]]:
    """Converts the results from a list of bundle inspector results into a dictionary of all of the array attributes

    The dictionary is useful for sorting as the key values will be unique.

    Args:
        inspected_values: List of bundle inspector outputs in the format produced by inspected_bundle_values()

    Returns:
        Dictionary of { Bundled Attribute Name : (TYPE, ROLE, IS_ARRAY, TUPLE_COUNT, VALUE) }
            TYPE: Attribute base type name
            ROLE: Attribute role name
            IS_ARRAY: Is the attribute an array type?
            TUPLE_COUNT: How many tuple elements in the attribute? (1 if simple data)
            VALUE: Value of the attribute extracted from the bundle
    """
    inspector_dictionary = {}
    for index in range(inspected_values[0]):
        inspector_dictionary[inspected_values[1][index]] = tuple(inspected_values[i][index] for i in range(2, 7))
    return inspector_dictionary


# ======================================================================
@ogt.deprecated_function("Use graph and/or node reference directly")
def graph_node_path(node_name: str) -> str:
    """Returns the full graph node prim path for the OmniGraph node of the given name in the test scene"""
    return f"/defaultPrim/graph/{node_name}"


# ======================================================================
@ogt.deprecated_function("Use og.Controller.create_prim()")
def _generate_test_file(
    node_definitions: List[Union[PrimNodeDefinitionType, OmniGraphNodeDefinitionType]], debug: bool = False
):
    """
    Generate a USDA file with data and connections that can exercise the bundle inspection node.
    The file could have be put into the data/ directory but it is easier to understand and modify here,
    especially since it is only useful in this test.

    Normally you wouldn't call this directly, you'd use set_up_test_stage() to do the next steps to load the scene.

    Args:
        node_definitions: List of USD node definitions for nodes in the test, not including the bundle source
            if the type is OmniGraphNodeDefinitionType then the values are tuples with OmniGraph node information:
                (NODE_NAME, NODE_TYPE, ATTRIBUTE_LINES, ATTRIBUTE_INFO)
                NODE_NAME: Local name of node in the graph
                NODE_TYPE: OmniGraph node type name
                ATTRIBUTE_LINES: List of the node's attribute declarations in USD form, including connections
                ATTRIBUTE_INFO: List of (TYPE:NAME) for attributes in the node
            else if the type is PrimNodeDefinitionType it contains data in the format for _write_usd_prim_definition
        debug: If True then alll bundle inspector nodes are put into debug mode, printing their computed results

    Returns:
        Path to the generated USDA file. The caller is responsible for deleting the file when done.
    """
    temp_fd = NamedTemporaryFile(mode="w+", suffix=".usda", delete=False)  # noqa: PLR1732
    out = ogt.IndentedOutput(temp_fd)

    # Start with a basic scene containing two cubes and a gather node
    out.write("#usda 1.0")
    out.write("(")
    out.write('    doc = """Bundle Inspection Node Test"""')
    out.write(")")
    out.write('def Xform "defaultPrim"')
    if out.indent("{"):
        for node_data in node_definitions:
            if len(node_data) == 2:
                _write_usd_prim_definition(node_data, out)

        out.write('def ComputeGraph "graph"')
        if out.indent("{"):
            out.write('def ComputeGraphSettings "computegraphSettings"')
            if out.indent("{"):
                out.write('custom token evaluator:type = "dirty_push"')
                out.exdent("}")

            for node_data in node_definitions:
                if len(node_data) != 4:
                    continue
                (node_name, node_type, attribute_lines, _) = node_data
                out.write(f'def ComputeNode "{node_name}"')
                if out.indent("{"):
                    out.write(f'custom token node:type = "{node_type}"')
                    out.write("custom int node:typeVersion = 1")
                    for attribute_line in attribute_lines:
                        assign = " = true" if debug and attribute_line.find("inputs:print") >= 0 else ""
                        out.write(f"{attribute_line}{assign}")
                    out.exdent("}")
            out.exdent("}")
        out.exdent("}")

    return temp_fd.name


# ======================================================================
@ogt.deprecated_function("Use og.Controller.edit() to set up a test scene")
async def set_up_test_stage(
    node_definitions: List[Union[PrimNodeDefinitionType, OmniGraphNodeDefinitionType]], debug: bool = False
):
    """Set up an initial stage based on a specific node type and its attributes

    Args:
        node_definitions: List of (NODE_NAME, NODE_TYPE_NAME, LIST_OF_USD_ATTRIBUTE_DECLARATIONS, ATTRIBUTE_INFO)
        debug: If True then all bundle inspector nodes are put into debug mode, printing their computed results

    Returns:
        Dictionary by name of a list of "Attribute" objects corresponding to the named attributes
    """
    path_to_file = _generate_test_file(node_definitions, debug)
    (result, error) = await ogts.load_test_file(path_to_file)
    if not result:
        raise ValueError(f"{error} on {path_to_file}")
    attributes = {}
    for node_data in node_definitions:
        node_name = node_data[0]
        test_node = og.get_current_graph().get_node(graph_node_path(node_name))
        if len(node_data) == 4:
            attribute_info = node_data[3]
            attributes[node_name] = [
                test_node.get_attribute(name) for attr_type, name in attribute_info if attr_type != "bundle"
            ]
    await get_app().next_update_async()
    return attributes


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.verify_bundles_are_equal()")
def compare_bundle_contents(inspector_attributes: List, bundle_inspector_output: Union[Dict[str, Any], List[Any]]):
    """Helper to compare the contents of two bundled attribute sets

    Args:
        inspector_attributes: Attribute list for the bundle inspector node
        bundle_inspector_output: Outputs retrieved from the inspector in the format generated by
                                 inspected_bundle_values() if it's a list, or the format generated by
                                 bundle_inspector_values_as_dictionary() if it's a dictionary

    Raises:
        ValueError if any of the bundle output values do not match the expected ones
    """
    # Mutating to a dictionary makes it easier to compare lists of lists
    if isinstance(bundle_inspector_output, list):
        bundled_size = bundle_inspector_output[0]
        bundled_data = bundle_inspector_values_as_dictionary(bundle_inspector_output)
    else:
        # If the data is already in dictionary form no conversion is needed
        bundled_size = len(bundle_inspector_output)
        bundled_data = bundle_inspector_output

    bundle_inspector_values = computed_bundle_values(inspector_attributes)
    if bundled_size != bundle_inspector_values[0]:
        raise ValueError(f"Bundled Count expected {bundled_size}, got {bundle_inspector_values[0]}")
    computed_values = bundle_inspector_values_as_dictionary(bundle_inspector_values)

    for attribute_name, attribute_values in computed_values.items():
        expected_attribute_values = bundled_data[attribute_name]
        for attribute_index, computed_value in enumerate(attribute_values):
            # Bundled attribute types not yet supported by the bundle inspector are silently ignored
            if isinstance(computed_value, (list, tuple)):
                if "__unsupported__" in computed_value:
                    continue
            else:
                if computed_value == "__unsupported__":
                    continue
            error_msg = f"Unexpected value on attribute {inspector_attributes[attribute_index + 2].get_name()}"
            ogts.verify_values(expected_attribute_values[attribute_index], computed_value, error_msg)


# ======================================================================
@ogt.deprecated_function("Use omni.graph.nodes.test.verify_bundles_are_equal()")
def compare_bundle_inspector_results(
    expected: BundleInspectorOutputType, actual: BundleInspectorOutputType, test: Optional[TestCase] = None
):
    """Helper to compare the expected outputs from a bundle inspector node

    This uses the unit test framework for the comparisons.

    Args:
        expected: Values expected to be retrieved from the bundle inspector's outputs
        actual: Outputs retrieved from the inspector
        test: If not None then use it to report errors found, otherwise use a locally constructed TestCase

    Raises:
        ValueError if any of the bundle output values do not match the expected ones
    """
    if test is None:
        test = TestCase()

    # Breaking open the tuples helps to clarify the operation being performed
    (
        expected_count,
        expected_names,
        expected_types,
        expected_roles,
        expected_array_depths,
        expected_tuple_counts,
        expected_values,
    ) = expected
    (
        actual_count,
        actual_names,
        actual_types,
        actual_roles,
        actual_array_depths,
        actual_tuple_counts,
        actual_values,
    ) = actual
    test.assertEqual(actual_count, expected_count)
    test.assertCountEqual(actual_names, expected_names)
    test.assertCountEqual(actual_types, expected_types)
    test.assertCountEqual(actual_roles, expected_roles)
    test.assertCountEqual(actual_array_depths, expected_array_depths)
    test.assertCountEqual(actual_tuple_counts, expected_tuple_counts)
    for one_actual, one_expected in zip(actual_values, expected_values):
        if isinstance(one_actual, list):
            test.assertCountEqual(one_actual, one_expected)
        else:
            test.assertEqual(one_actual, one_expected)


# ==============================================================================================================
@ogt.deprecated_function("Use og.Controller.create_prim or og.Controller.edit for creating prims")
def create_prim_with_everything(prim_path: Optional[str] = None) -> Tuple[Usd.Prim, Dict[str, Any]]:
    """Deprecated function; should no longer be called"""
    return (None, {})
