lib

class DirmapCache:
    """Caching class to get settings and sitesync easily and only once."""
    _project_name = None
    _project_settings = None
    _sitesync_addon_discovered = False
    _sitesync_addon = None
    _mapping = None

    @classmethod
    def project_name(cls):
        if cls._project_name is None:
            cls._project_name = os.getenv("AYON_PROJECT_NAME")
        return cls._project_name

    @classmethod
    def project_settings(cls):
        if cls._project_settings is None:
            cls._project_settings = get_project_settings(cls.project_name())
        return cls._project_settings

    @classmethod
    def sitesync_addon(cls):
        if not cls._sitesync_addon_discovered:
            cls._sitesync_addon_discovered = True
            cls._sitesync_addon = AddonsManager().get("sitesync")
        return cls._sitesync_addon

    @classmethod
    def mapping(cls):
        return cls._mapping

    @classmethod
    def set_mapping(cls, mapping):
        cls._mapping = mapping

class Knobby(object):
    """[DEPRECATED] For creating knob which it's type isn't
                    mapped in `create_knobs`

    Args:
        type (string): Nuke knob type name
        value: Value to be set with `Knob.setValue`, put `None` if not required
        flags (list, optional): Knob flags to be set with `Knob.setFlag`
        *args: Args other than knob name for initializing knob class

    """

    def __init__(self, type, value, flags=None, *args):
        self.type = type
        self.value = value
        self.flags = flags or []
        self.args = args

    def create(self, name, nice=None):
        knob_cls = getattr(nuke, self.type)
        knob = knob_cls(name, nice, *self.args)
        if self.value is not None:
            knob.setValue(self.value)
        for flag in self.flags:
            knob.setFlag(flag)
        return knob

    @staticmethod
    def nice_naming(key):
        """Convert camelCase name into UI Display Name"""
        words = re.findall('[A-Z][^A-Z]*', key[0].upper() + key[1:])
        return " ".join(words)

class NukeDirmap(HostDirmap):
    def __init__(self, file_name, *args, **kwargs):
        """
        Args:
            file_name (str): full path of referenced file from workfiles
            *args (tuple): Positional arguments for 'HostDirmap' class
            **kwargs (dict): Keyword arguments for 'HostDirmap' class
        """

        self.file_name = file_name
        super(NukeDirmap, self).__init__(*args, **kwargs)

    def on_enable_dirmap(self):
        pass

    def dirmap_routine(self, source_path, destination_path):
        source_path = source_path.lower().replace(os.sep, '/')
        destination_path = destination_path.lower().replace(os.sep, '/')
        if platform.system().lower() == "windows":
            self.file_name = self.file_name.lower().replace(
                source_path, destination_path)
        else:
            self.file_name = self.file_name.replace(
                source_path, destination_path)

class WorkfileSettings(object):
    """
    All settings for workfile will be set

    This object is setting all possible root settings to the workfile.
    Including Colorspace, Frame ranges, Resolution format. It can set it
    to Root node or to any given node.

    Arguments:
        root (node): nuke's root node
        nodes (list): list of nuke's nodes
        nodes_filter (list): filtering classes for nodes

    """

    def __init__(self, root_node=None, nodes=None, **kwargs):
        project_entity = kwargs.get("project")
        if project_entity is None:
            project_name = get_current_project_name()
            project_entity = ayon_api.get_project(project_name)
        else:
            project_name = project_entity["name"]

        Context._project_entity = project_entity
        self._project_name = project_name
        self._folder_path = get_current_folder_path()
        self._folder_entity = ayon_api.get_folder_by_path(
            project_name, self._folder_path
        )
        self._task_name = get_current_task_name()
        self._context_label = "{} > {}".format(self._folder_path,
                                               self._task_name)
        self._task_entity = ayon_api.get_task_by_name(
            project_name,
            self._folder_entity["id"],
            self._task_name
        )
        self._root_node = root_node or nuke.root()
        self._nodes = self.get_nodes(nodes=nodes)

        context_data = get_template_data_with_names(
            project_name, self._folder_path, self._task_name, "nuke"
        )
        self.formatting_data = context_data

    def get_nodes(self, nodes=None, nodes_filter=None):

        if not isinstance(nodes, list) and not isinstance(nodes_filter, list):
            return [n for n in nuke.allNodes()]
        elif not isinstance(nodes, list) and isinstance(nodes_filter, list):
            nodes = list()
            for filter in nodes_filter:
                [nodes.append(n) for n in nuke.allNodes(filter=filter)]
            return nodes
        elif isinstance(nodes, list) and not isinstance(nodes_filter, list):
            return [n for n in self._nodes]
        elif isinstance(nodes, list) and isinstance(nodes_filter, list):
            for filter in nodes_filter:
                return [n for n in self._nodes if filter in n.Class()]

    # TODO: move into ./colorspace.py
    def set_viewers_colorspace(self, imageio_nuke):
        ''' Adds correct colorspace to viewer

        Arguments:
            imageio_nuke (dict): nuke colorspace configurations

        '''
        filter_knobs = [
            "viewerProcess",
            "wipe_position",
            "monitorOutOutputTransform"
        ]
        viewer_process = get_formatted_display_and_view(
            imageio_nuke["viewer"], self.formatting_data, self._root_node
        )
        output_transform = get_formatted_display_and_view(
            imageio_nuke["monitor"], self.formatting_data, self._root_node
        )
        erased_viewers = []
        for v in nuke.allNodes(filter="Viewer"):
            # set viewProcess to preset from settings
            v["viewerProcess"].setValue(viewer_process)

            if viewer_process not in v["viewerProcess"].value():
                copy_inputs = v.dependencies()
                copy_knobs = {
                    k: v[k].value() for k in v.knobs()
                    if k not in filter_knobs
                }

                # delete viewer with wrong settings
                erased_viewers.append(v["name"].value())
                nuke.delete(v)

                # create new viewer
                nv = nuke.createNode("Viewer")

                # connect to original inputs
                for i, n in enumerate(copy_inputs):
                    nv.setInput(i, n)

                # set copied knobs
                for knob_name, knob_value in copy_knobs.items():
                    nv[knob_name].setValue(knob_value)

                # set viewerProcess
                nv["viewerProcess"].setValue(viewer_process)
                nv["monitorOutOutputTransform"].setValue(output_transform)

        if erased_viewers:
            log.warning(
                "Attention! Viewer nodes {} were erased."
                "It had wrong color profile".format(erased_viewers))

    # TODO: move into ./colorspace.py
    def set_root_colorspace(self, imageio_host):
        ''' Adds correct colorspace to root

        Arguments:
            imageio_host (dict): host colorspace configurations

        '''
        config_data = get_current_context_imageio_config_preset()

        workfile_settings = imageio_host["workfile"]
        color_management = workfile_settings["color_management"]
        native_ocio_config = workfile_settings["native_ocio_config"]

        if not config_data:
            # no ocio config found and no custom path used
            if self._root_node["colorManagement"].value() \
                        not in color_management:
                self._root_node["colorManagement"].setValue(color_management)

            # second set ocio version
            if self._root_node["OCIO_config"].value() \
                        not in native_ocio_config:
                self._root_node["OCIO_config"].setValue(native_ocio_config)

        else:
            # OCIO config path is defined from prelaunch hook
            self._root_node["colorManagement"].setValue("OCIO")

            # print previous settings in case some were found in workfile
            residual_path = self._root_node["customOCIOConfigPath"].value()
            if residual_path:
                log.info("Residual OCIO config path found: `{}`".format(
                    residual_path
                ))

        # set ocio config path
        if config_data:
            config_path = config_data["path"].replace("\\", "/")
            log.info("OCIO config path found: `{}`".format(
                config_path))

            # check if there's a mismatch between environment and settings
            correct_settings = self._is_settings_matching_environment(
                config_data)

            # if there's no mismatch between environment and settings
            if correct_settings:
                self._set_ocio_config_path_to_workfile(config_data)

        workfile_settings_output = {}
        # get monitor lut from settings respecting Nuke version differences
        monitor_lut_data = self._get_monitor_settings(
            workfile_settings["monitor_out_lut"],
            workfile_settings["monitor_lut"]
        )
        workfile_settings_output.update(monitor_lut_data)
        workfile_settings_output.update(
            {
                "workingSpaceLUT": workfile_settings["working_space"],
                "int8Lut": workfile_settings["int_8_lut"],
                "int16Lut": workfile_settings["int_16_lut"],
                "logLut": workfile_settings["log_lut"],
                "floatLut": workfile_settings["float_lut"],
            }
        )

        # then set the rest
        for knob, value_ in workfile_settings_output.items():
            # skip unfilled ocio config path
            # it will be dict in value
            if isinstance(value_, dict):
                continue
            # skip empty values
            if not value_:
                continue
            self._root_node[knob].setValue(str(value_))

    def _get_monitor_settings(self, viewer_lut, monitor_lut):
        """ Get monitor settings from viewer and monitor lut

        Args:
            viewer_lut (str): viewer lut string
            monitor_lut (str): monitor lut string

        Returns:
            dict: monitor settings
        """
        output_data = {}
        m_display, m_viewer = get_viewer_config_from_string(monitor_lut)
        v_display, v_viewer = get_viewer_config_from_string(viewer_lut)

        # set monitor lut differently for nuke version 14
        if nuke.NUKE_VERSION_MAJOR >= 14:
            output_data["monitorOutLUT"] = create_viewer_profile_string(
                m_viewer, m_display, path_like=False)
            # monitorLut=thumbnails - viewerProcess makes more sense
            output_data["monitorLut"] = create_viewer_profile_string(
                v_viewer, v_display, path_like=False)

        if nuke.NUKE_VERSION_MAJOR == 13:
            output_data["monitorOutLUT"] = create_viewer_profile_string(
                m_viewer, m_display, path_like=False)
            # monitorLut=thumbnails - viewerProcess makes more sense
            output_data["monitorLut"] = create_viewer_profile_string(
                v_viewer, v_display, path_like=True)
        if nuke.NUKE_VERSION_MAJOR <= 12:
            output_data["monitorLut"] = create_viewer_profile_string(
                m_viewer, m_display, path_like=True)

        return output_data

    def _is_settings_matching_environment(self, config_data):
        """ Check if OCIO config path is different from environment

        Args:
            config_data (dict): OCIO config data from settings

        Returns:
            bool: True if settings are matching environment, False otherwise
        """
        current_ocio_path = os.environ["OCIO"]
        settings_ocio_path = config_data["path"]

        # normalize all paths to forward slashes
        current_ocio_path = current_ocio_path.replace("\\", "/")
        settings_ocio_path = settings_ocio_path.replace("\\", "/")

        if current_ocio_path != settings_ocio_path:
            message = """
It seems like there's a mismatch between the OCIO config path set in your Nuke
settings and the actual path set in your OCIO environment.

To resolve this, please follow these steps:
1. Close Nuke if it's currently open.
2. Reopen Nuke.

Please note the paths for your reference:

- The OCIO environment path currently set:
  `{env_path}`

- The path in your current Nuke settings:
  `{settings_path}`

Reopening Nuke should synchronize these paths and resolve any discrepancies.
"""
            nuke.message(
                message.format(
                    env_path=current_ocio_path,
                    settings_path=settings_ocio_path
                )
            )
            return False

        return True

    def _set_ocio_config_path_to_workfile(self, config_data):
        """ Set OCIO config path to workfile

        Path set into nuke workfile. It is trying to replace path with
        environment variable if possible. If not, it will set it as it is.
        It also saves the script to apply the change, but only if it's not
        empty Untitled script.

        Args:
            config_data (dict): OCIO config data from settings

        """
        # replace path with env var if possible
        ocio_path = self._replace_ocio_path_with_env_var(config_data)

        log.info("Setting OCIO config path to: `{}`".format(
            ocio_path))

        self._root_node["customOCIOConfigPath"].setValue(
            ocio_path
        )
        self._root_node["OCIO_config"].setValue("custom")

        # only save script if it's not empty
        if self._root_node["name"].value() != "":
            log.info("Saving script to apply OCIO config path change.")
            nuke.scriptSave()

    def _get_included_vars(self, config_template):
        """ Get all environment variables included in template

        Args:
            config_template (str): OCIO config template from settings

        Returns:
            list: list of environment variables included in template
        """
        # resolve all environments for whitelist variables
        included_vars = [
            "BUILTIN_OCIO_ROOT",
        ]

        # include all project root related env vars
        for env_var in os.environ:
            if env_var.startswith("AYON_PROJECT_ROOT_"):
                included_vars.append(env_var)

        # use regex to find env var in template with format {ENV_VAR}
        # this way we make sure only template used env vars are included
        env_var_regex = r"\{([A-Z0-9_]+)\}"
        env_var = re.findall(env_var_regex, config_template)
        if env_var:
            included_vars.append(env_var[0])

        return included_vars

    def _replace_ocio_path_with_env_var(self, config_data):
        """ Replace OCIO config path with environment variable

        Environment variable is added as TCL expression to path. TCL expression
        is also replacing backward slashes found in path for windows
        formatted values.

        Args:
            config_data (str): OCIO config dict from settings

        Returns:
            str: OCIO config path with environment variable TCL expression
        """
        config_path = config_data["path"].replace("\\", "/")
        config_template = config_data["template"]

        included_vars = self._get_included_vars(config_template)

        # make sure we return original path if no env var is included
        new_path = config_path

        for env_var in included_vars:
            env_path = os.getenv(env_var)
            if not env_path:
                continue

            # it has to be directory current process can see
            if not os.path.isdir(env_path):
                continue

            # make sure paths are in same format
            env_path = env_path.replace("\\", "/")
            path = config_path.replace("\\", "/")

            # check if env_path is in path and replace to first found positive
            if env_path in path:
                # with regsub we make sure path format of slashes is correct
                resub_expr = (
                    "[regsub -all {{\\\\}} [getenv {}] \"/\"]").format(env_var)

                new_path = path.replace(
                    env_path, resub_expr
                )
                break

        return new_path

    # TODO: move into ./colorspace.py
    def set_writes_colorspace(self):
        ''' Adds correct colorspace to write node dict

        '''
        for node in nuke.allNodes(filter="Group", group=self._root_node):
            log.info("Setting colorspace to `{}`".format(node.name()))

            # get data from avalon knob
            avalon_knob_data = read_avalon_data(node)
            node_data = get_node_data(node, INSTANCE_DATA_KNOB)

            if (
                # backward compatibility
                # TODO: remove this once old avalon data api will be removed
                avalon_knob_data
                and avalon_knob_data.get("id") not in {
                    AYON_INSTANCE_ID, AVALON_INSTANCE_ID
                }
            ):
                continue
            elif (
                node_data
                and node_data.get("id") not in {
                    AYON_INSTANCE_ID, AVALON_INSTANCE_ID
                }
            ):
                continue

            if (
                # backward compatibility
                # TODO: remove this once old avalon data api will be removed
                avalon_knob_data
                and "creator" not in avalon_knob_data
            ):
                continue
            elif (
                node_data
                and "creator_identifier" not in node_data
            ):
                continue

            nuke_imageio_writes = None
            if avalon_knob_data:
                # establish families
                product_type = avalon_knob_data.get("productType")
                if product_type is None:
                    product_type = avalon_knob_data["family"]
                families = [product_type]
                if avalon_knob_data.get("families"):
                    families.append(avalon_knob_data.get("families"))

                nuke_imageio_writes = get_imageio_node_setting(
                    node_class=avalon_knob_data["families"],
                    plugin_name=avalon_knob_data["creator"],
                    product_name=avalon_knob_data["productName"]
                )
            elif node_data:
                nuke_imageio_writes = get_write_node_template_attr(node)

            if not nuke_imageio_writes:
                return

            write_node = None

            # get into the group node
            node.begin()
            for x in nuke.allNodes():
                if x.Class() == "Write":
                    write_node = x
            node.end()

            if not write_node:
                return

            set_node_knobs_from_settings(
                write_node, nuke_imageio_writes["knobs"])

    # TODO: move into ./colorspace.py
    def set_reads_colorspace(self, read_clrs_inputs):
        """ Setting colorspace to Read nodes

        Looping through all read nodes and tries to set colorspace based
        on regex rules in presets
        """
        changes = {}
        for n in nuke.allNodes():
            file = nuke.filename(n)
            if n.Class() != "Read":
                continue

            # check if any colorspace presets for read is matching
            preset_clrsp = None

            for input in read_clrs_inputs:
                if not bool(re.search(input["regex"], file)):
                    continue
                preset_clrsp = input["colorspace"]

            if preset_clrsp is not None:
                current = n["colorspace"].value()
                future = str(preset_clrsp)
                if current != future:
                    changes[n.name()] = {
                        "from": current,
                        "to": future
                    }

        if changes:
            msg = "Read nodes are not set to correct colorspace:\n\n"
            for nname, knobs in changes.items():
                msg += (
                    " - node: '{0}' is now '{1}' but should be '{2}'\n"
                ).format(nname, knobs["from"], knobs["to"])

            msg += "\nWould you like to change it?"

            if nuke.ask(msg):
                for nname, knobs in changes.items():
                    n = nuke.toNode(nname)
                    n["colorspace"].setValue(knobs["to"])
                    log.info(
                        "Setting `{0}` to `{1}`".format(
                            nname,
                            knobs["to"]))

    # TODO: move into ./colorspace.py
    def set_colorspace(self):
        ''' Setting colorspace following presets
        '''
        # get imageio
        nuke_colorspace = get_nuke_imageio_settings()

        log.info("Setting colorspace to workfile...")
        try:
            self.set_root_colorspace(nuke_colorspace)
        except AttributeError as _error:
            msg = "Set Colorspace to workfile error: {}".format(_error)
            nuke.message(msg)

        log.info("Setting colorspace to viewers...")
        try:
            self.set_viewers_colorspace(nuke_colorspace)
        except AttributeError as _error:
            msg = "Set Colorspace to viewer error: {}".format(_error)
            nuke.message(msg)

        log.info("Setting colorspace to write nodes...")
        try:
            self.set_writes_colorspace()
        except AttributeError as _error:
            nuke.message(_error)
            log.error(_error)

        log.info("Setting colorspace to read nodes...")
        read_clrs_inputs = nuke_colorspace["regex_inputs"].get("inputs", [])
        if read_clrs_inputs:
            self.set_reads_colorspace(read_clrs_inputs)

    def reset_frame_range_handles(self):
        """Set frame range to current folder."""

        if "attrib" not in self._task_entity:
            msg = "Task {} doesn't have set any 'attrib'".format(
                self._context_label
            )
            log.warning(msg)
            nuke.message(msg)
            return

        task_attributes = self._task_entity["attrib"]

        missing_cols = []
        check_cols = ["fps", "frameStart", "frameEnd",
                      "handleStart", "handleEnd"]

        for col in check_cols:
            if col not in task_attributes:
                missing_cols.append(col)

        if len(missing_cols) > 0:
            missing = ", ".join(missing_cols)
            msg = "'{}' are not set for task '{}'!".format(
                missing, self._context_label)
            log.warning(msg)
            nuke.message(msg)
            return

        # get handles values
        handle_start = task_attributes["handleStart"]
        handle_end = task_attributes["handleEnd"]
        frame_start = task_attributes["frameStart"]
        frame_end = task_attributes["frameEnd"]

        fps = float(task_attributes["fps"])
        frame_start_handle = frame_start - handle_start
        frame_end_handle = frame_end + handle_end

        self._root_node["lock_range"].setValue(False)
        self._root_node["fps"].setValue(fps)
        self._root_node["first_frame"].setValue(frame_start_handle)
        self._root_node["last_frame"].setValue(frame_end_handle)
        self._root_node["lock_range"].setValue(True)

        # update node graph so knobs are updated
        update_node_graph()

        frame_range = '{0}-{1}'.format(frame_start, frame_end)

        for node in nuke.allNodes(filter="Viewer"):
            node['frame_range'].setValue(frame_range)
            node['frame_range_lock'].setValue(True)
            node['frame_range'].setValue(frame_range)
            node['frame_range_lock'].setValue(True)

        if not ASSIST:
            set_node_data(
                self._root_node,
                INSTANCE_DATA_KNOB,
                {
                    "handleStart": int(handle_start),
                    "handleEnd": int(handle_end)
                }
            )
        else:
            log.warning(
                "NukeAssist mode is not allowing "
                "updating custom knobs..."
            )

    def reset_resolution(self):
        """Set resolution to project resolution."""
        log.info("Resetting resolution")
        project_name = get_current_project_name()
        task_attributes = self._task_entity["attrib"]

        format_data = {
            "width": task_attributes["resolutionWidth"],
            "height": task_attributes["resolutionHeight"],
            "pixel_aspect": task_attributes["pixelAspect"],
            "name": project_name
        }

        if any(x_ for x_ in format_data.values() if x_ is None):
            msg = ("Missing set shot attributes in DB."
                   "\nContact your supervisor!."
                   "\n\nWidth: `{width}`"
                   "\nHeight: `{height}`"
                   "\nPixel Aspect: `{pixel_aspect}`").format(**format_data)
            log.error(msg)
            nuke.message(msg)

        existing_format = None
        for format in nuke.formats():
            if format_data["name"] == format.name():
                existing_format = format
                break

        if existing_format:
            # Enforce existing format to be correct.
            existing_format.setWidth(format_data["width"])
            existing_format.setHeight(format_data["height"])
            existing_format.setPixelAspect(format_data["pixel_aspect"])
        else:
            format_string = self.make_format_string(**format_data)
            log.info("Creating new format: {}".format(format_string))
            nuke.addFormat(format_string)

        nuke.root()["format"].setValue(format_data["name"])
        log.info("Format is set.")

        # update node graph so knobs are updated
        update_node_graph()

    def make_format_string(self, **kwargs):
        if kwargs.get("r"):
            return (
                "{width} "
                "{height} "
                "{x} "
                "{y} "
                "{r} "
                "{t} "
                "{pixel_aspect:.2f} "
                "{name}".format(**kwargs)
            )
        else:
            return (
                "{width} "
                "{height} "
                "{pixel_aspect:.2f} "
                "{name}".format(**kwargs)
            )

    def set_context_settings(self):
        # replace reset resolution from avalon core to pype's
        self.reset_resolution()
        # replace reset resolution from avalon core to pype's
        self.reset_frame_range_handles()
        # add colorspace menu item
        self.set_colorspace()

    def set_favorites(self):
        from .utils import set_context_favorites

        work_dir = os.getenv("AYON_WORKDIR")
        # TODO validate functionality
        # - does expect the structure is '{root}/{project}/{folder}'
        # - this used asset name expecting it is unique in project
        folder_path = get_current_folder_path()
        folder_name = folder_path.split("/")[-1]
        favorite_items = OrderedDict()

        # project
        # get project's root and split to parts
        projects_root = os.path.normpath(work_dir.split(
            Context.project_name)[0])
        # add project name
        project_dir = os.path.join(projects_root, Context.project_name) + "/"
        # add to favorites
        favorite_items.update({"Project dir": project_dir.replace("\\", "/")})

        # folder
        folder_root = os.path.normpath(work_dir.split(
            folder_name)[0])
        # add folder name
        folder_dir = os.path.join(folder_root, folder_name) + "/"
        # add to favorites
        favorite_items.update({"Shot dir": folder_dir.replace("\\", "/")})

        # workdir
        favorite_items.update({"Work dir": work_dir.replace("\\", "/")})

        set_context_favorites(favorite_items)