"""Extension implementation required by Carbonite extension loader"""
# These are used by the extension and require explicit import here
import importlib
import json
import time
from contextlib import suppress
from pathlib import Path
from typing import List, Set, Tuple, Union

import carb
import carb.profiler
import omni.ext
import omni.graph.core as og
import omni.graph.tools as ogt
from omni.kit.app import get_app_interface
from omni.kit.commands import unregister_module_commands
from pxr import Sdf

from .v1_5_0.update_file_format import (
    cb_update_to_include_schema,
    migrate_attribute_custom_data,
    remove_global_compute_graph,
)


# ==============================================================================================================
def register_categories():
    """If the Kit location is available then use it to find category configuration files and register their contents"""
    try:  # noqa: PLR1702
        config_dir = Path(carb.tokens.get_tokens_interface().resolve("${kit}")) / "dev" / "ogn" / "config"
        if config_dir.is_dir():
            node_categories = og.get_node_categories_interface()

            for config_file in config_dir.iterdir():
                try:
                    with open(config_file, "r", encoding="utf-8") as json_fd:
                        configuration = json.load(json_fd)
                    if "categoryDefinitions" in configuration:
                        for category_name, category_description in configuration["categoryDefinitions"].items():
                            if category_name[0] == "$":
                                continue
                            node_categories.define_category(category_name, category_description)
                except json.decoder.JSONDecodeError:
                    pass
    except Exception as error:  # noqa: PLW0703
        carb.log_warn(f"Could not find the OmniGraph configuration directory - {error}")


# ==============================================================================================================
class _PublicExtension(omni.ext.IExt):
    """Mandatory extension instantiation required by the Carbonite extension handler"""

    # Globally accessible timing information tracking how long it takes to register and deregister nodes in an extension
    REGISTRATION_TIMING = {}
    DEREGISTRATION_TIMING = {}

    def __init__(self, *args, **kwargs):
        self.__extension_disabled_hook = None
        self.__extension_enabled_hook = None
        self.__interface = None
        self.__migrate_attribute_custom_data_handle = None
        self.__module_registrations = {}
        self.__pre_del_prim_cb = None
        self.__pre_del_prim_cb = None
        self.__remove_global_graph_handle = None
        self.__schema_upgrade_handle = None

    def on_after_ext_enabled(self, ext_id: str, *_):
        ogt.dbg_reg("Checking for Python nodes in {}", ext_id)
        start_registration = time.perf_counter_ns()
        carb.profiler.begin(1, f"OmniGraph.PythonNode.Registration.{ext_id}")
        manager = get_app_interface().get_extension_manager()
        with suppress(KeyError):
            for module_info in manager.get_extension_dict(ext_id).get_dict()["python"]["module"]:
                module_name = module_info["name"]
                with suppress(ModuleNotFoundError):
                    ogn_module = importlib.import_module(f"{module_name}.ogn")
                    self.__module_registrations.setdefault(ext_id, []).append(
                        og.PythonNodeRegistration(ogn_module, module_name)
                    )
                    ogt.dbg_reg(" -> Registered nodes from module {}", module_name)
        carb.profiler.end(1)
        end_registration = time.perf_counter_ns()
        self.REGISTRATION_TIMING[ext_id] = end_registration - start_registration

    def on_before_ext_disabled(self, ext_id: str, *_):
        start_registration = time.perf_counter_ns()
        carb.profiler.begin(1, f"OmniGraph.PythonNode.Deregistration.{ext_id}")
        with suppress(KeyError):
            del self.__module_registrations[ext_id]
            ogt.dbg_reg("Deregistered Python nodes in {}", ext_id)
        carb.profiler.end(1)
        end_registration = time.perf_counter_ns()
        self.DEREGISTRATION_TIMING[ext_id] = end_registration - start_registration

    # --------------------------------------------------------------------------------------------------------------
    # begin-file-format-update
    def on_startup(self):
        """Set up initial conditions for the Python part of the extension"""
        self.__module_registrations = {}
        self.__interface = og.acquire_interface()
        og.register_python_node()

        hooks = omni.kit.app.get_app_interface().get_extension_manager().get_hooks()
        self.__extension_enabled_hook = hooks.create_extension_state_change_hook(
            self.on_after_ext_enabled, omni.ext.ExtensionStateChangeType.AFTER_EXTENSION_ENABLE
        )
        assert self.__extension_enabled_hook
        self.__extension_disabled_hook = hooks.create_extension_state_change_hook(
            self.on_before_ext_disabled, omni.ext.ExtensionStateChangeType.BEFORE_EXTENSION_DISABLE
        )
        assert self.__extension_disabled_hook
        self.__remove_global_graph_handle = og.register_pre_load_file_format_upgrade_callback(
            remove_global_compute_graph
        )
        self.__migrate_attribute_custom_data_handle = og.register_pre_load_file_format_upgrade_callback(
            migrate_attribute_custom_data
        )
        self.__schema_upgrade_handle = og.register_pre_load_file_format_upgrade_callback(cb_update_to_include_schema)

        self.__pre_del_prim_cb = omni.kit.commands.register_callback(
            "DeletePrimsCommand", omni.kit.commands.PRE_DO_CALLBACK, self.__on_before_del_prim
        )

        # Categories are in configuration files so they have to be added to OmniGraph here
        register_categories()

    # --------------------------------------------------------------------------------------------------------------
    def on_shutdown(self):
        """Shutting down this part of the extension prepares it for hot reload"""
        omni.kit.commands.unregister_callback(self.__pre_del_prim_cb)
        self.__pre_del_prim_cb = None
        og.on_shutdown()
        og.release_interface(self.__interface)
        unregister_module_commands(og.cmds)
        og.cmds = None
        self.__extension_enabled_hook = None
        self.__extension_disabled_hook = None
        if self.__module_registrations:
            carb.log_warn(f"Shutting down OmniGraph with modules still registered - {self.__module_registrations}")

        og.deregister_pre_load_file_format_upgrade_callback(self.__remove_global_graph_handle)
        self.__remove_global_graph_handle = None
        og.deregister_pre_load_file_format_upgrade_callback(self.__migrate_attribute_custom_data_handle)
        self.__migrate_attribute_custom_data_handle = None
        og.deregister_pre_load_file_format_upgrade_callback(self.__schema_upgrade_handle)
        self.__schema_upgrade_handle = None

    # end-file-format-update

    # --------------------------------------------------------------------------------------------------------------
    def __on_before_del_prim(self, info):
        # Get the command's 'paths' argument.
        paths: List[Union[str, Sdf.Path]] = info.get("paths", None)
        if not paths:
            return

        # Remove duplicates.
        paths = set(paths)

        # Find all the OG nodes which will be removed.
        nodes_being_removed: Set[og.Node] = set()
        while paths:
            path = paths.pop()
            graph = og.Controller.graph(path)
            if graph:
                # Add the subgraph's nodes.
                for node in graph.get_nodes():
                    paths.add(node.get_prim_path())
            else:
                with suppress(og.OmniGraphValueError):
                    node = og.Controller.node(path)
                    nodes_being_removed.add(node)

        # Find any connections where the source node is being deleted
        # but the destination isn't. Those will result in broken connections
        # which the OG core won't be able to restore if the DeletePrims
        # command is undone. To work around that we'll add an undo task
        # to the command's undo block.
        conns_to_undo: List[Tuple[str, str]] = []
        for node in nodes_being_removed:
            for attr in node.get_attributes():
                for dest_attr in attr.get_downstream_connections():
                    dest_node = dest_attr.get_node()
                    if dest_node not in nodes_being_removed:
                        # Attribute.get_path() does not return a valid path for bundle outputs so
                        # we must build the path from its parts.
                        src_path = node.get_prim_path() + "." + attr.get_name()
                        dest_path = dest_node.get_prim_path() + "." + dest_attr.get_name()
                        conns_to_undo.append((src_path, dest_path))
        if conns_to_undo:
            omni.kit.commands.execute("_OGRestoreConnectionsOnUndo", connections=conns_to_undo)
