# Copyright (c) 2018-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.
#
__all__ = ["StageItem"]

import carb
import omni.ui as ui
import omni.kit.usd.layers as layers
import omni.usd

from .models import PrimNameModel, TypeModel, VisibilityModel
from pxr import Sdf, UsdGeom, Usd
from typing import List


class StageItem(ui.AbstractItem):
    """
    A single AbstractItemModel item that represents a single prim.
    StageItem is a cached view of the prim.
    """

    def __init__(
        self,
        path: Sdf.Path,
        stage,
        stage_model,
        # Deprecated
        flat=False,
        root_identifier=None,
        load_payloads=False,
        check_missing_references=False,
    ):
        super().__init__()
        self.__stage_model = stage_model

        # Internal access
        self._path = path

        # Prim handle
        self.__prim = None

        # Filtering
        self.filtered = None
        # True if it has any descendant that is filtered
        self.child_filtered = None

        # defaultPrim
        self.__is_default = False

        # True if it has missing references
        self.__missing_references = None

        # True if it has references authored
        self.__has_references = False

        # All references/payloads for this prim
        self.__payrefs = set()

        # If this prim includes references/payloads that are out of sync.
        self.__is_outdated = False

        # True if it has payloads authored
        self.__has_payloads = False

        # True if it's visible
        self.__visible = True

        # True if is in session
        self.__in_session = False

        # True if it's instanceable
        self.__instanceable = False

        # True if this item should be auto-loaded when its
        # references or payloads are outdated.
        self.__auto_reload = False

        # True if it's children of instance.
        self.__instance_proxy = False

        self.__display_name = None

        # Lazy load
        self.__flags_updated = True

        # Models for builtin columns. Those models
        # will only be accessed internally.
        self.__name_model = None
        self.__type_model = None
        self.__visibility_model = None

    def destroy(self):
        self.__stage = None
        self.__stage_model = None
        self.__payrefs.clear()

    @property
    def path(self):
        """Prim path."""

        return self._path

    @property
    def stage_model(self):
        """StageModel this StageItem belongs to."""

        return self.__stage_model

    @property
    def usd_context(self) -> omni.usd.UsdContext:
        return self.__stage_model.usd_context if self.__stage_model else None

    @property
    def stage(self) -> Usd.Stage:
        """USD stage the prim belongs to."""

        return self.__stage_model.stage if self.__stage_model else None

    @property
    def payrefs(self) -> List[str]:
        """All external references and payloads that influence this prim."""

        self.__update_flags_internal()

        return list(self.__payrefs)

    @property
    def is_default(self) -> bool:
        """Whether this prim is the default prim or not."""

        self.__update_flags_internal()

        return self.__is_default

    @property
    def is_outdated(self) -> bool:
        """Whether this prim includes references or payloads that has new changes to fetch or not."""

        self.__update_flags_internal()
        return self.__is_outdated

    @property
    def in_session(self) -> bool:
        """Whether this prim includes references or payloads that are in session or not."""

        self.__update_flags_internal()

        return self.__in_session

    @property
    def auto_reload(self):
        """Whether this prim is configured as auto-reload when its references or payloads are outdated."""

        self.__update_flags_internal()

        return self.__auto_reload

    @property
    def root_identifier(self) -> str:
        """Returns the root layer's identifier of the stage this prim belongs to."""

        stage = self.stage
        return stage.GetRootLayer().identifier if stage else None

    @property
    def instance_proxy(self) -> bool:
        """Whether the prim is an instance proxy or not."""

        self.__update_flags_internal()

        return self.__instance_proxy

    @property
    def instanceable(self) -> bool:
        """True when it has `instanceable` flag enabled"""

        self.__update_flags_internal()

        return self.__instanceable

    @property
    def visible(self) -> bool:
        self.__update_flags_internal()

        return self.__visible

    @property
    def payloads(self):
        """True when the prim has authored payloads"""

        self.__update_flags_internal()

        return self.__has_payloads

    @property
    def references(self) -> bool:
        """True when the prim has authored references"""

        self.__update_flags_internal()

        return self.__has_references

    @property
    def name(self) -> str:
        """The path name."""

        return self._path.name

    @property
    def display_name(self) -> str:
        """The display name of prim inside the metadata."""

        self.__update_flags_internal()

        return self.__display_name or self.name

    @property
    def prim(self):
        """The prim handle."""

        self.__update_flags_internal()

        return self.__prim

    @property
    def active(self):
        """True if the prim is active."""

        self.__update_flags_internal()

        if not self.__prim:
            return False

        return self.__prim.IsActive()

    @property
    def type_name(self):
        """Type name of the prim."""

        self.__update_flags_internal()

        prim = self.__prim
        if not prim:
            return ""

        return prim.GetTypeName()

    @property
    def has_missing_references(self):
        """
        Whether the prim includes any missing references or payloads or not.
        It checkes the references/payloads recursively.
        """

        self.__update_flags_internal()

        return self.__missing_references

    @property
    def children(self):
        """Returns children items. If children are not populated yet, they will be populated."""

        return self.__stage_model.get_item_children(self)

    def update_flags(self, prim=None):
        """Refreshes item states from USD."""

        # Lazy load.
        self.__flags_updated = True

    def __update_prim(self):
        stage = self.stage
        self.__prim = stage.GetPrimAtPath(self._path) if stage else None

    def __update_visibility(self):
        prim = self.prim
        if prim and prim.IsA(UsdGeom.Imageable):
            visibility = UsdGeom.Imageable(prim).ComputeVisibility()
            self.__visible = visibility != UsdGeom.Tokens.invisible
        else:
            self.__visible = False

    def __get_payrefs(self):
        if not self.__has_references and not self.__has_payloads:
            return set()

        result = set()
        references = omni.usd.get_composed_references_from_prim(self.__prim)
        payloads = omni.usd.get_composed_payloads_from_prim(self.__prim)

        for reference, layer in payloads + references:
            if not reference.assetPath:
                continue

            absolute_path = layer.ComputeAbsolutePath(reference.assetPath)
            result.add(absolute_path)

        return result

    def __get_live_status(self):
        if not self.__stage_model.usd_context:
            return

        if self.__payrefs:
            live_syncing = layers.get_live_syncing(self.__stage_model.usd_context)
            for absolute_path in self.__payrefs:
                if live_syncing.is_prim_in_live_session(self.prim.GetPath(), absolute_path):
                    return True

        return False

    def __get_outdated_status(self):
        if not self.__stage_model.usd_context:
            return

        if self.__payrefs:
            layers_state_interface = layers.get_layers_state(self.__stage_model.usd_context)
            for absolute_path in self.__payrefs:
                if layers_state_interface.is_layer_outdated(absolute_path):
                    return True

        return False

    def __get_auto_reload_status(self):
        if not self.__stage_model.usd_context:
            return

        if carb.settings.get_settings().get_as_bool(layers.SETTINGS_AUTO_RELOAD_NON_SUBLAYERS):
            if not self.is_outdated:
                return True

        if self.__payrefs:
            layers_state_interface = layers.get_layers_state(self.__stage_model.usd_context)
            for absolute_path in self.__payrefs:
                if layers_state_interface.is_auto_reload_layer(absolute_path):
                    return True

        return False

    def __has_missing_payrefs(self):
        prim = self.prim
        # If prim is not loaded.
        if prim and not prim.IsLoaded():
            return False

        for identifier in self.payrefs:
            layer = Sdf.Find(identifier)
            if not layer:
                return True

        return False

    def __update_flags_internal(self):
        # Param is for back compitability.
        if not self.__flags_updated:
            return

        self.__flags_updated = False

        self.__update_prim()
        prim = self.prim
        if not prim:
            return

        self.__is_default = prim == self.stage.GetDefaultPrim()
        self.__display_name = omni.usd.editor.get_display_name(prim) or ""
        # Refresh model to decide which name to display.
        if self.__name_model:
            self.__name_model.rebuild()
        self.__update_visibility()
        self.__has_references = prim.HasAuthoredReferences()
        self.__has_payloads = prim.HasAuthoredPayloads()
        self.__payrefs = self.__get_payrefs()
        self.__instanceable = prim.IsInstanceable()
        self.__instance_proxy = prim.IsInstanceProxy()
        self.__missing_references = self.__has_missing_payrefs()
        self.__is_outdated = self.__get_outdated_status()
        self.__in_session = self.__get_live_status()
        self.__auto_reload = self.__get_auto_reload_status()

    def __repr__(self):
        return f"<Omni::UI Stage Item '{self._path}'>"

    def __str__(self):
        return f"{self._path}"

    # Internal properties
    @property
    def name_model(self):
        self.__update_flags_internal()

        if not self.__name_model:
            self.__name_model = PrimNameModel(self)

        return self.__name_model

    @property
    def type_model(self):
        self.__update_flags_internal()

        if not self.__type_model:
            self.__type_model = TypeModel(self)

        return self.__type_model

    @property
    def visibility_model(self):
        self.__update_flags_internal()

        if not self.__visibility_model:
            self.__visibility_model = VisibilityModel(self)

        return self.__visibility_model

    @property
    def is_flat(self):
        return self.__stage_model.flat if self.__stage_model else False

    @is_flat.setter
    def is_flat(self, flat: bool):
        # Rebuild it only when it's populated
        if self.__name_model:
            self.__name_model.rebuild()

    @property
    def load_payloads(self):
        return self.stage_model.load_payloads if self.stage_model else None

    # Deprecated functions
    @property
    def check_missing_references(self):
        """Deprecated: It will always check missing references now."""

        return True

    @check_missing_references.setter
    def check_missing_references(self, value):
        """Deprecated."""
        pass

    def set_default_prim(self, is_default):
        """Deprecated."""
        pass
