import carb
import omni.ext
import omni.usd

from typing import List, Union, Dict
from .layers_interface import Layers, LiveSyncing, LayersState, AutoAuthoring
from .._omni_kit_usd_layers import acquire_layers_interface, release_layers_interface, LayerEditMode
from pxr import Sdf, Usd

_all_layers_instances: Dict[omni.usd.UsdContext, Layers] = {}
_layers_interface = None


def get_layers_interface():
    global _layers_interface
    if not _layers_interface:
        _layers_interface = acquire_layers_interface()

    return _layers_interface


def get_layers(context_name_or_instance: Union[str, omni.usd.UsdContext] = "") -> Union[Layers, None]:
    global _all_layers_instances
    if not context_name_or_instance:
        context_name_or_instance = ""

    if isinstance(context_name_or_instance, str):
        usd_context = omni.usd.get_context(context_name_or_instance)
    elif isinstance(context_name_or_instance, omni.usd.UsdContext):
        usd_context = context_name_or_instance
    else:
        carb.log_warn("Failed to get layers instance since the param must be name or instance of UsdContext.")
        return None

    if not usd_context:
        carb.log_warn("Failed to query layers instance since UsdContext cannot be found.")
        return None

    layers_instance = _all_layers_instances.get(usd_context, None)
    if not layers_instance:
        instance_native = get_layers_interface().get_layers_instance_by_context(usd_context)
        if instance_native:
            layers_instance = Layers(instance_native, usd_context)
            _all_layers_instances[usd_context] = layers_instance
        else:
            layers_instance = None
            carb.log_warn("Failed to query layers instance.")

    return layers_instance


def get_auto_authoring(context_name_or_instance: Union[str, omni.usd.UsdContext] = "") -> Union[AutoAuthoring, None]:
    layers = get_layers(context_name_or_instance)
    if layers:
        return layers.get_auto_authoring()
    
    return None


def get_layers_state(context_name_or_instance: Union[str, omni.usd.UsdContext] = "") -> Union[LayersState, None]:
    layers = get_layers(context_name_or_instance)
    if layers:
        return layers.get_layers_state()
    
    return None


def get_live_syncing(context_name_or_instance: Union[str, omni.usd.UsdContext] = "") -> Union[LiveSyncing, None]:
    layers = get_layers(context_name_or_instance)
    if layers:
        return layers.get_live_syncing()
    
    return None


def active_authoring_layer_context(usd_context) -> Usd.EditContext:
    """
    Gets the edit context for edit target if it's in non-auto authoring mode,
    or edit context for default edit layer if it's in auto authoring mode.
    """
    stage = usd_context.get_stage()
    layers = get_layers(usd_context)
    auto_authoring = layers.get_auto_authoring()
    edit_mode = layers.get_edit_mode()
    if edit_mode == LayerEditMode.NORMAL:
        edit_target_layer = stage.GetEditTarget().GetLayer()
    else:
        default_edit_layer_identifier = auto_authoring.get_default_layer()
        edit_target_layer = Sdf.Find(default_edit_layer_identifier)

    if not edit_target_layer:
        edit_target_layer = stage.GetRootLayer()

    return Usd.EditContext(stage, edit_target_layer)


# Use extension entry points to acquire and release interface.
class Extension(omni.ext.IExt):
    def on_startup(self):
        # Load layers for default context
        get_layers()

    def on_shutdown(self):
        global _all_layers_instances

        for _, instance in _all_layers_instances.items():
            instance.destroy()
        _all_layers_instances.clear()

        global _layers_interface
        if _layers_interface:
            release_layers_interface(_layers_interface)
