"""File format upgrade utilities"""
from typing import Any, Dict, List, Optional, Tuple

import carb
import omni.graph.core as og
import omni.usd
from omni.kit.commands import execute
from pxr import Gf, OmniGraphSchema, Usd

from .utils import ALLOW_IMPLICIT_GRAPH_SETTING, USE_SCHEMA_PRIMS_SETTING

EXTENDED_ATTRIBUTE_SCOPED_CUSTOM_DATA_VERSION = og.FileFormatVersion(1, 5)

LAST_FILE_FORMAT_VERSION_WITHOUT_SCHEMA = og.FileFormatVersion(1, 10)
"""TODO: This is the last file format version where it is possible to have OmniGraph prims that do not use the schema
         It should be updated when the final deprecation is made.
"""


# ==============================================================================================================
def remove_global_compute_graph(old_format_version, new_format_version, graph):
    """If the setting to allow a global implicit graph is not on then remove it if found
    and replace it with a regular graph"""

    settings = carb.settings.get_settings()
    allow_implicit_graph = settings.get(ALLOW_IMPLICIT_GRAPH_SETTING)
    # If the schema is enabled it will fix the global compute graphs too
    use_schema_prims = settings.get(og.USE_SCHEMA_PRIMS_SETTING)

    if not allow_implicit_graph and not use_schema_prims:
        stage = omni.usd.get_context().get_stage()
        if stage is not None:
            iterator = iter(stage.TraverseAll())
            for prim in iterator:
                if prim.GetTypeName() == "GlobalComputeGraph":
                    carb.log_info(
                        "Converting " + prim.GetPrimPath().pathString + "from GlobalComputeGraph to ComputeGraph"
                    )
                    prim.SetTypeName("ComputeGraph")
                    iterator.PruneChildren()


# ==============================================================================================================
def check_for_bare_connections(stage: Usd.Stage) -> List[str]:
    """Check for any of the old unsupported connections from nodes to prims that cannot exist in schema-based world

    Return:
        List of strings describing the forbidden connections found
    """
    iterator = iter(stage.TraverseAll())
    bad_connections = []
    for prim in iterator:
        if prim.GetTypeName() in ["ComputeNode", "OmniGraphNode"]:
            for attribute in prim.GetAttributes():
                for connected_path in attribute.GetConnections():
                    connected_prim = stage.GetPrimAtPath(connected_path.GetPrimPath())
                    if not connected_prim.IsValid():
                        carb.log_warn(
                            f"Invalid connection found at {attribute.GetPath()}:"
                            f" {connected_path.GetPrimPath()} while migrating to new schema"
                        )
                    elif connected_prim.GetTypeName() not in ["ComputeNode", "OmniGraphNode"]:
                        bad_connections.append(f"{prim.GetPrimPath()} - {connected_path}")
    return bad_connections


# ==============================================================================================================
def migrate_attribute_custom_data(
    old_format_version: og.FileFormatVersion, new_format_version: og.FileFormatVersion, _
):
    """Add scoped customData to attributes which have the old format"""
    if old_format_version < EXTENDED_ATTRIBUTE_SCOPED_CUSTOM_DATA_VERSION:
        stage = omni.usd.get_context().get_stage()
        if stage is None:
            return
        iterator = iter(stage.TraverseAll())
        for prim in iterator:
            if prim.GetTypeName() in ["OmniGraphNode", "ComputeNode"]:
                for attribute in prim.GetAttributes():
                    custom_data = attribute.GetCustomData()
                    extended_type = custom_data.get("ExtendedAttributeType", None)
                    if extended_type:
                        attribute.SetCustomDataByKey("omni:graph:attrType", extended_type)
                iterator.PruneChildren()


# ==============================================================================================================
class OmniGraphSchemaMigrator:
    """Collection of methods to manage all of the steps involved in migrating an arbitrary graph to use schema prims"""

    def __init__(self):
        """Set up initial values"""
        self.__command_count = 0

    # --------------------------------------------------------------------------------------------------------------
    @property
    def stage(self) -> Usd.Stage:
        """Returns the stage in use - set as a property so that the latest version is always used as it changes"""
        return omni.usd.get_context().get_stage()

    # --------------------------------------------------------------------------------------------------------------
    def __check_for_old_prim_nodes(self) -> List[str]:
        """Check for any of the old unsupported prim nodes that cannot exist in schema-based world

        Return:
            List of prim paths for any unsupported prim node types
        """
        iterator = iter(self.stage.TraverseAll())
        prim_nodes = []
        for prim in iterator:
            if prim.GetTypeName() == "ComputeNode":
                node_type_attr = prim.GetAttribute("node:type")
                if node_type_attr.IsValid() and node_type_attr.Get() == "omni.graph.core.Prim":
                    prim_nodes.append(prim.GetPrimPath())
        return prim_nodes

    # --------------------------------------------------------------------------------------------------------------
    def __change_prim_types_to_match_schema(self) -> List[Tuple[str, str, str]]:
        """Modify any old prim names to match the ones used by the schema

        Return:
            List of (path, old_type, new_type) for prims whose type was changed
        """
        iterator = iter(self.stage.TraverseAll())
        retyped = []
        for prim in iterator:
            if prim.GetTypeName() in ["GlobalComputeGraph", "ComputeGraph"]:
                carb.log_info(f"Converting graph {prim.GetPrimPath()} from {prim.GetTypeName()} to OmniGraph")
                retyped.append((str(prim.GetPrimPath()), prim.GetTypeName(), "OmniGraph"))
                prim.SetTypeName("OmniGraph")
            elif prim.GetTypeName() == "ComputeNode":
                carb.log_info(f"Converting node {prim.GetPrimPath()} from ComputeNode to OmniGraphNode")
                retyped.append((str(prim.GetPrimPath()), prim.GetTypeName(), "OmniGraphNode"))
                prim.SetTypeName("OmniGraphNode")
                # Ensure that the schema-defined attributes on the node are no longer custom
                node_prim = OmniGraphSchema.OmniGraphNode(prim)
                if node_prim:
                    node_type_attr = node_prim.GetNodeTypeAttr()
                    if node_type_attr.IsValid():
                        if not node_type_attr.SetCustom(False):
                            carb.log_warn(f"Failed to set node:type attribute to non-custom on {prim.GetPrimPath()}")
                    else:
                        carb.log_warn(f"OmniGraph prim node {prim.GetPrimPath()} has no node:type attribute")
                    node_type_version_attr = node_prim.GetNodeTypeVersionAttr()
                    if node_type_version_attr.IsValid():
                        if not node_type_version_attr.SetCustom(False):
                            carb.log_warn(
                                f"Failed to set node:typeVersion attribute to non-custom on {prim.GetPrimPath()}"
                            )
                    else:
                        carb.log_warn(f"OmniGraph prim node {prim.GetPrimPath()} has no node:type attribute")
                else:
                    # This might be an error or it might just be the fact that the USD notification to change the type
                    # hasn't been processed yet. As this is only temporary compatibility code and we know that when
                    # reading a file the prims are immediately valid there is no need to add the complexity here to
                    # tie into an event loop and await the USD notification propagation.
                    pass
        return retyped

    # --------------------------------------------------------------------------------------------------------------
    def __create_top_level_graph(self) -> List[Tuple[str, str]]:
        """If there are any nodes or settings not in a graph then create a graph for them and move them into it.

        Assumes that before calling this all of the old prim type names have been replaced with the schema prim
        type names.

        Args:
            stage: USD stage on which to replace settings

        Return:
            List of (old_path, new_path) for prims that were moved from the root level to a new graph

        Raises:
            og.OmniGraphError if the required prim changes failed
        """
        iterator = iter(self.stage.TraverseAll())
        # Collect this map of settings prim path onto the path at which the parent graph must be created
        prims_to_move = []
        for prim in iterator:
            if prim.GetTypeName() in ["ComputeGraphSettings", "OmniGraphNode"]:
                parent = prim.GetParent()
                if not parent.IsValid() or parent.GetTypeName() != "OmniGraph":
                    prims_to_move.append(prim)
            elif prim.GetTypeName() in ["ComputeGraph", "OmniGraph"]:
                iterator.PruneChildren()
        prim_paths_moved = []
        if prims_to_move:
            carb.log_info(f"Creating top level graph for global nodes {[prim.GetPrimPath() for prim in prims_to_move]}")
            for prim in prims_to_move:
                parent_path = str(prim.GetParent().GetPrimPath())
                if parent_path[-1] != "/":
                    parent_path += "/"
                default_global_graph_path = f"{parent_path}__graphUsingSchemas"
                global_graph_prim = self.stage.GetPrimAtPath(default_global_graph_path)
                if not global_graph_prim:
                    (status, result) = execute("CreatePrim", prim_path=default_global_graph_path, prim_type="OmniGraph")
                    if not status:
                        raise og.OmniGraphError(
                            f"Error creating OmniGraph prim at {default_global_graph_path} - {result}"
                        )
                new_prim_path = f"{default_global_graph_path}/{prim.GetName()}"
                prim_paths_moved.append((str(prim.GetPrimPath()), new_prim_path))
                carb.log_info(f"Move {prim.GetPrimPath()} -> {new_prim_path}")
                (status, result) = execute("MovePrim", path_from=prim.GetPrimPath(), path_to=new_prim_path)
                if not status:
                    raise og.OmniGraphError(
                        f"Error moving OmniGraph prim from {prim.GetPrimPath()} to {new_prim_path} - {result}"
                    )
        return prim_paths_moved

    # --------------------------------------------------------------------------------------------------------------
    def __replace_settings_with_properties(self, new_file_format_version: Tuple[int, int]) -> List[str]:
        """If there are any of the old settings prims move their property values into their containing graph.

        Args:
            new_file_format_version: The file format version that should be used for the graph setting

        Return:
            List of prim paths with settings that were transferred to the containing graph and then deleted
        """
        iterator = iter(self.stage.TraverseAll())
        prims_to_remove = []
        for prim in iterator:  # noqa: PLR1702
            if prim.GetTypeName() == "ComputeGraphSettings":
                carb.log_info(f"Moving settings from {prim.GetPrimPath()} into the parent graph")
                graph_prim = prim.GetParent()
                # Need to check against the old types as well as the USD change notices may not have percolated
                if graph_prim.IsValid() and graph_prim.GetTypeName() in [
                    "GlobalComputeGraph",
                    "ComputeGraph",
                    "OmniGraph",
                ]:
                    for attr in prim.GetAttributes():
                        setting_name = attr.GetName()
                        setting_value = attr.Get()
                        if setting_name == "flatCacheBacking":
                            setting_name = "fabricCacheBacking"
                            if setting_value == "StagedWithHistory":
                                setting_value = "StageWithHistory"
                        graph_attr = graph_prim.GetAttribute(setting_name)
                        if not graph_attr.IsValid():
                            graph_attr = graph_prim.CreateAttribute(
                                setting_name, attr.GetTypeName(), custom=False, variability=attr.GetVariability()
                            )
                        if graph_attr.IsValid():
                            graph_attr.Set(setting_value)
                        else:
                            carb.log_warn(
                                f"Could not create settings attribute {attr.GetName()}"
                                f" on graph {graph_prim.GetPrimPath()}"
                            )
                    prims_to_remove.append(str(prim.GetPrimPath()))
                    schema_prim = OmniGraphSchema.OmniGraph(graph_prim)
                    if bool(schema_prim):
                        file_format_version_attr = schema_prim.GetFileFormatVersionAttr()
                        file_format_version_attr.Set(Gf.Vec2i(new_file_format_version))
                    else:
                        carb.log_warn(f"Could not cast graph prim {graph_prim.GetPrimPath()} to OmniGraph schema")
                else:
                    carb.log_warn(f"Could not find graph above {prim.GetPrimPath()} to receive settings")
            elif prim.GetTypeName() == "ComputeNodeMetadata":
                # Take advantage of the loop to also get rid of the obsolete metadata children
                prims_to_remove.append(str(prim.GetPrimPath()))
        if prims_to_remove:
            (status, result) = execute("DeletePrims", paths=prims_to_remove)
            if not status:
                raise og.OmniGraphError(f"Error deleting prims {prims_to_remove} - {result}")
        return prims_to_remove

    # --------------------------------------------------------------------------------------------------------------
    # Handling for old files that did not use the OmniGraph schema (earlier than 1.3)
    def update(self, new_file_format_version: Optional[Tuple[int, int]] = None) -> Dict[str, Any]:
        """Update the current file to use the new schema.

        This can run either before OmniGraph has attached to the stage (e.g. when reading in an old file) or when an
        explicit migration has been requested (e.g. when an old file exists and the user wants to migrate it). In both
        cases if any prim types have been migrated the useSchemaPrims setting will be forced to True so that subsequent
        OmniGraph additions use the schema types. (Mostly because this is easier than trying to track the state of
        whether a graph contains such prims already.)

        Conversion of old scenes entails:
            - Changing any ComputeGraph or GlobalComputeGraph prims to be OmniGraph types
            - If a ComputeGraphSettings prim exists, migrating its attribute values to the OmniGraph prim
            - Changing any ComputeNode prims to be OmniGraphNode types
            - If any ComputeNode prims appear without a parent ComputeGraph then create a default graph
              and move them into it

        Args:
            new_file_format_version: The file format version that should be used for the graph setting. If None then
                                     it will force the version immediately following the schema conversion (1, 4)

        Return:
            A dictionary of operations that were performed as part of the migration (key describes the operation, value
            is the objects to which it was applied)
        """
        stage = omni.usd.get_context().get_stage()
        if stage is None:
            return {"message": "No stage to convert"}

        try:
            self.__command_count = 0
            return_values = {}

            # Verify that the old prim nodes don't exist
            old_prim_nodes = self.__check_for_old_prim_nodes()
            if old_prim_nodes:
                raise og.OmniGraphError(
                    f"Deprecated omni.graph.core.Prim nodes not allowed with schema - {old_prim_nodes}"
                )

            # Verify that bare connections from OmniGraph nodes to non-OmniGraph prims do not exist
            bare_connections = check_for_bare_connections(self.stage)
            if bare_connections:
                raise og.OmniGraphError(
                    "Deprecated connections from an OmniGraph node to a USD Prim not allowed with schema"
                    f" - {bare_connections}"
                )

            # First pass - change ComputeGraph/GlobalComputeGraph -> OmniGraph and ComputeNode -> OmniGraphNode
            return_values["Prim Types Changed"] = self.__change_prim_types_to_match_schema()

            # Second pass - if there are any OmniGraphNode or ComputeGraphSettings prims not in a graph,
            # create one and move them into it
            return_values["Root Prims Moved To Graph"] = self.__create_top_level_graph()

            # Third pass - move settings from their own prim to the parent graph prim
            if new_file_format_version is None:
                new_file_format_version = (1, 4)
            return_values["Settings Removed"] = self.__replace_settings_with_properties(new_file_format_version)

            # Now that the scene has been migrated the schema setting has to be enabled or bad things will happen.
            # Don't bother doing it if nothing changed though.
            if any(return_values.values()):
                carb.settings.get_settings().set(USE_SCHEMA_PRIMS_SETTING, True)

            return return_values
        except og.OmniGraphError as error:
            # If anything failed it would leave the graph in a hybrid state so try to restore it back to its original
            # state so that it remains stable.
            carb.log_error(str(error))
            for _i in range(self.__command_count):
                omni.kit.undo()
            return_values = {}

        return return_values


# ==============================================================================================================
# Handling for old files that did not use the OmniGraph schema (earlier than 1.3)
def update_to_include_schema(new_file_format_version: Optional[Tuple[int, int]] = None) -> Dict[str, Any]:
    """Update the current file to use the new schema.  See OmniGraphSchemaMigrator.update() for details

    This can run either before OmniGraph has attached to the stage (e.g. when reading in an old file) or when an
    explicit migration has been requested (e.g. when an old file exists and the user wants to migrate it). In both
    cases if any prim types have been migrated the useSchemaPrims setting will be forced to True so that subsequent
    OmniGraph additions use the schema types. (Mostly because this is easier than trying to track the state of whether
    a graph contains such prims already.)

    Conversion of old scenes entails:
        - Changing any ComputeGraph or GlobalComputeGraph prims to be OmniGraph types
        - If a ComputeGraphSettings prim exists, migrating its attribute values to the OmniGraph prim
        - Changing any ComputeNode prims to be OmniGraphNode types
        - If any ComputeNode prims appear without a parent ComputeGraph create a default graph and move them into it

    Args:
        new_file_format_version: The file format version that should be used for the graph setting. If None then it will
                                 force the version immediately following the schema conversion (1, 4)

    Return:
        A dictionary of operations that were performed as part of the migration (key describes the operation, value
        is the objects to which it was applied)

    Raises:
        og.OmniGraphError if any of the attempted changes failed - will attempt to restore graph to original state
    """
    return OmniGraphSchemaMigrator().update()


# ==============================================================================================================
# Handling for old files that did not use the OmniGraph schema (earlier than 1.3)
def cb_update_to_include_schema(
    old_version: Optional[og.FileFormatVersion], new_version: Optional[og.FileFormatVersion], graph: Optional[og.Graph]
) -> Dict[str, Any]:
    """Callback invoked when a file is loaded to update old files to use the new schema.

    This will be called anytime a file is loaded with a non-current version. The old version and new version are
    checked to confirm that the values cross over the boundary when schemas were created, and if so then the schema
    information is applied to the file.

    Args:
        old_version: Version the file to upgrade uses
        new_version: Current file version expected
        graph: Graph to convert (only present for historical reasons - the entire stage is updated)

    Return:
        A dictionary of operations that were performed as part of the migration (key describes the operation, value
        is the objects to which it was applied)
    """
    # If the setting to automatically migrate is not on then there is nothing to do here
    if not carb.settings.get_settings().get(USE_SCHEMA_PRIMS_SETTING):
        return {}

    # If the file format version is one of the ones that must contain schema prims no migration is needed
    if old_version is not None and (
        old_version.majorVersion > LAST_FILE_FORMAT_VERSION_WITHOUT_SCHEMA.majorVersion
        or (
            old_version.majorVersion == LAST_FILE_FORMAT_VERSION_WITHOUT_SCHEMA.majorVersion
            and old_version.minorVersion > LAST_FILE_FORMAT_VERSION_WITHOUT_SCHEMA.minorVersion
        )
    ):
        return {"message": f"File format version {old_version} already uses the schema"}

    return update_to_include_schema((new_version.majorVersion, new_version.minorVersion))


# --------------------------------------------------------------------------------------------------------------
RESETTING_USE_SCHEMA_PRIMS = False
"""Global variable to avoid potential infinite recursion when restoring the useSchemaPrims setting"""
RESETTING_IMPLICIT_GRAPH = False
"""Global variable to avoid potential infinite recursion when restoring the allowGlobalImplicitGraph setting"""


# --------------------------------------------------------------------------------------------------------------
def can_set_use_schema_prims_setting(new_value: bool) -> str:
    """Check to see if a new value for the schema prims setting is valid - returns warning text, empty if okay"""

    def __prim_types_in_scene(prim_types: List[str]) -> bool:
        """Returns true if the scene contains prims whose types are any in the list"""
        found_prims = []
        stage = omni.usd.get_context().get_stage()
        if stage is not None:
            iterator = iter(stage.TraverseAll())
            for prim in iterator:
                if prim.GetTypeName() in prim_types:
                    found_prims.append(f"{prim.GetPrimPath()} : {prim.GetTypeName()}")
        return found_prims

    if new_value:
        forbidden_prims = __prim_types_in_scene(
            ["ComputeGraphSettings", "ComputeNode", "ComputeGraph", "GlobalComputeGraph"]
        )
        if forbidden_prims:
            return f"Cannot enable {USE_SCHEMA_PRIMS_SETTING} when the scene contains old style prims {forbidden_prims}"
        bare_connections = check_for_bare_connections(check_for_bare_connections(omni.usd.get_context().get_stage()))
        if bare_connections:
            return (
                f"Cannot enable {USE_SCHEMA_PRIMS_SETTING} when the scene contains"
                f" direct prim connections {bare_connections}"
            )
    else:
        forbidden_prims = __prim_types_in_scene(["OmniGraph", "OmniGraphNode"])
        if forbidden_prims:
            return f"Cannot disable {USE_SCHEMA_PRIMS_SETTING} when the scene contains schema prims {forbidden_prims}"

    return ""
