# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
import asyncio
import omni.kit.commands
import omni.timeline
import omni.ui as ui
import omni.usd

import weakref
import omni.kit.widget.layers

from pxr import Usd, Sdf, UsdGeom, UsdPhysics
from omni.kit.widget.layers import LayerUtils
from .types import LayerMetaType


class LayerPathModel(ui.SimpleStringModel):
    def __init__(self, layer_item: weakref):
        super().__init__()
        self._layer_item = layer_item
        if self._layer_item and self._layer_item():
            self._identifier = self._layer_item().identifier
        else:
            self._identifier = None

    def get_value_as_string(self):
        return self._identifier

    def set_value(self, value):
        if value != self._identifier:
            self._identifier = value
            self._value_changed()

    def begin_edit(self):
        pass

    def end_edit(self):
        if not self._layer_item or not self._layer_item():
            return

        value = self.get_value_as_string()
        layer_item = self._layer_item()
        if not layer_item.reserved:
            sublayer_position = LayerUtils.get_sublayer_position_in_parent(
                layer_item.parent.identifier, layer_item.identifier
            )
            omni.kit.commands.execute(
                "ReplaceSublayer",
                layer_identifier=layer_item.parent.identifier,
                sublayer_position=sublayer_position,
                new_layer_path=value,
            )

        # OM-76598: Delay frames until layer item is initialized.
        async def focus_on_layer_item(path):
            await omni.kit.app.get_app().next_update_async()
            await omni.kit.app.get_app().next_update_async()
            layer_instance = omni.kit.widget.layers.get_instance()
            layer_instance.set_current_focused_layer_item(path)
        
        asyncio.ensure_future(focus_on_layer_item(value))

    def replace_layer(self, value):
        self.set_value(value)
        self.end_edit()

    def is_reserved_layer(self):
        if not self._layer_item or not self._layer_item():
            return True

        return self._layer_item().reserved

    def anonymous(self):
        if not self._layer_item or not self._layer_item() or not self._layer_item().layer:
            return True

        return self._layer_item().anonymous


class LayerWorldAxisItem(ui.AbstractItem):
    def __init__(self, text):
        super().__init__()
        self.model = ui.SimpleStringModel(text)


class LayerWorldAxisModel(ui.AbstractItemModel):
    def __init__(self, layer_item: weakref):
        super().__init__()

        self._layer_item = layer_item
        self._items = [LayerWorldAxisItem(text) for text in [UsdGeom.Tokens.y, UsdGeom.Tokens.z]]
        self._current_index = ui.SimpleIntModel()
        self.on_value_changed()
        self._current_index.add_value_changed_fn(self._current_index_changed)

    def get_item_children(self, item):
        return self._items

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._current_index

        return item.model

    def _current_index_changed(self, model):
        if not self._layer_item or not self._layer_item():
            return

        stage = Usd.Stage.Open(self._layer_item().layer)
        index = model.as_int
        if index == 0:
            omni.kit.commands.execute("ModifyStageAxis", stage=stage, axis=UsdGeom.Tokens.y)
        elif index == 1:
            omni.kit.commands.execute("ModifyStageAxis", stage=stage, axis=UsdGeom.Tokens.z)
            UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)

        self._item_changed(None)

    def get_usd_token_name(self):
        return UsdGeom.Tokens.upAxis

    def on_value_changed(self):
        if self._layer_item and self._layer_item():
            layer = self._layer_item().layer
            stage = Usd.Stage.Open(layer)
            up_axis = UsdGeom.GetStageUpAxis(stage)
            index = 0 if up_axis == UsdGeom.Tokens.y else 1
            self._current_index.set_value(index)


class LayerMetaModel(ui.AbstractValueModel):
    def __init__(self, layer_item: weakref, meta_type: LayerMetaType):
        super().__init__()
        self._layer_item = layer_item
        self._meta_type = meta_type
        self._value = self._get_value_as_string()

    def get_value_as_string(self):
        return self._value

    def _get_value_as_string(self):
        if not self._layer_item or not self._layer_item():
            return None

        layer = self._layer_item().layer
        if not layer:
            return None

        if self._meta_type == LayerMetaType.COMMENT:
            return str(layer.comment)
        elif self._meta_type == LayerMetaType.DOC:
            return str(layer.documentation)
        elif self._meta_type == LayerMetaType.START_TIME:
            return str(layer.startTimeCode)
        elif self._meta_type == LayerMetaType.END_TIME:
            return str(layer.endTimeCode)
        elif self._meta_type == LayerMetaType.TIMECODES_PER_SECOND:
            return str(layer.timeCodesPerSecond)
        elif layer.HasFramesPerSecond() and self._meta_type == LayerMetaType.FPS_PER_SECOND:
            return str(layer.framesPerSecond)
        elif self._meta_type == LayerMetaType.UNITS:
            stage = Usd.Stage.Open(layer, None, None, Usd.Stage.LoadNone)
            meters = UsdGeom.GetStageMetersPerUnit(stage)
            return str(meters)
        elif self._meta_type == LayerMetaType.KG_PER_UNIT:
            stage = Usd.Stage.Open(layer, None, None, Usd.Stage.LoadNone)
            kilograms = UsdPhysics.GetStageKilogramsPerUnit(stage)
            return str(kilograms)
        elif (
            self._layer_item().parent and self._meta_type == LayerMetaType.LAYER_OFFSET
        ):
            parent = self._layer_item().parent
            layer_index = LayerUtils.get_sublayer_position_in_parent(parent.identifier, layer.identifier)
            offset = parent.layer.subLayerOffsets[layer_index]
            return str(offset.offset)
        elif self._layer_item().parent and self._meta_type == LayerMetaType.LAYER_SCALE:
            parent = self._layer_item().parent
            layer_index = LayerUtils.get_sublayer_position_in_parent(parent.identifier, layer.identifier)
            offset = parent.layer.subLayerOffsets[layer_index]
            return str(offset.scale)

        return None

    def set_value(self, value):
        if value != self._value:
            self._value = value
            self._value_changed()

    def end_edit(self):
        if not self._layer_item or not self._layer_item():
            return

        layer = self._layer_item().layer
        if not layer:
            return

        if self._layer_item().parent:
            parent_layer_identifier = self._layer_item().parent.identifier
        else:
            parent_layer_identifier = None

        layer_identifier = layer.identifier
        omni.kit.commands.execute(
            "ModifyLayerMetadata",
            layer_identifier=layer_identifier,
            parent_layer_identifier=parent_layer_identifier,
            meta_index=self._meta_type,
            value=self._value,
        )

    def begin_edit(self):
        pass

    def get_usd_token_name(self):
        if self._meta_type == LayerMetaType.COMMENT:
            return Sdf.Layer.CommentKey
        elif self._meta_type == LayerMetaType.DOC:
            return Sdf.Layer.DocumentationKey
        elif self._meta_type == LayerMetaType.START_TIME:
            return Sdf.Layer.StartTimeCodeKey
        elif self._meta_type == LayerMetaType.END_TIME:
            return Sdf.Layer.EndTimeCodeKey
        elif self._meta_type == LayerMetaType.TIMECODES_PER_SECOND:
            return Sdf.Layer.TimeCodesPerSecondKey
        elif self._meta_type == LayerMetaType.FPS_PER_SECOND:
            return Sdf.Layer.FramesPerSecondKey
        elif self._meta_type == LayerMetaType.UNITS:
            return UsdGeom.Tokens.metersPerUnit
        elif self._meta_type == LayerMetaType.KG_PER_UNIT:
            return UsdPhysics.Tokens.kilogramsPerUnit
        elif self._meta_type == LayerMetaType.LAYER_OFFSET:
            return "subLayerOffsets_offset"  # USD has no python bindings for this key
        elif self._meta_type == LayerMetaType.LAYER_SCALE:
            return "subLayerOffsets_scale"  # USD has no python bindings for this key

        return ""

    def on_value_changed(self):
        self.set_value(self._get_value_as_string())
