lib

def get_channel_format(stack_path, channel):
    """Retrieve the channel format of a specific stack channel.

    See `alg.mapexport.channelFormat` (javascript API) for more details.

    The channel format data is:
        "label" (str): The channel format label: could be one of
            [sRGB8, L8, RGB8, L16, RGB16, L16F, RGB16F, L32F, RGB32F]
        "color" (bool): True if the format is in color, False is grayscale
        "floating" (bool): True if the format uses floating point
            representation, false otherwise
        "bitDepth" (int): Bit per color channel (could be 8, 16 or 32 bpc)

    Arguments:
        stack_path (list or str): Path to the stack, could be
            "Texture set name" or ["Texture set name", "Stack name"]
        channel (str): Identifier of the channel to export
            (see `get_channel_identifiers`)

    Returns:
        dict: The channel format data.

    """
    stack_path = _convert_stack_path_to_cmd_str(stack_path)
    cmd = f"alg.mapexport.channelFormat({stack_path}, '{channel}')"
    return substance_painter.js.evaluate(cmd)

def get_channel_identifiers(stack_path=None):
    """Return the list of channel identifiers.

    If a context is passed (texture set/stack),
    return only used channels with resolved user channels.

    Channel identifiers are:
        basecolor, height, specular, opacity, emissive, displacement,
        glossiness, roughness, anisotropylevel, anisotropyangle, transmissive,
        scattering, reflection, ior, metallic, normal, ambientOcclusion,
        diffuse, specularlevel, blendingmask, [custom user names].

    Args:
        stack_path (list or str, Optional): Path to the stack, could be
            "Texture set name" or ["Texture set name", "Stack name"]

    Returns:
        list: List of channel identifiers.

    """
    if stack_path is None:
        stack_path = ""
    else:
        stack_path = _convert_stack_path_to_cmd_str(stack_path)
    cmd = f"alg.mapexport.channelIdentifiers({stack_path})"
    return substance_painter.js.evaluate(cmd)

def get_document_structure():
    """Dump the document structure.

    See `alg.mapexport.documentStructure` (javascript API) for more details.

    Returns:
        dict: Document structure or None when no project is open

    """
    return substance_painter.js.evaluate("alg.mapexport.documentStructure()")

def get_export_presets():
    """Return Export Preset resource URLs for all available Export Presets.

    Returns:
        dict: {Resource url: GUI Label}

    """
    # TODO: Find more optimal way to find all export templates

    preset_resources = {}
    for shelf in substance_painter.resource.Shelves.all():
        shelf_path = os.path.normpath(shelf.path())

        presets_path = os.path.join(shelf_path, "export-presets")
        if not os.path.exists(presets_path):
            continue

        for filename in os.listdir(presets_path):
            if filename.endswith(".spexp"):
                template_name = os.path.splitext(filename)[0]

                resource = substance_painter.resource.ResourceID(
                    context=shelf.name(),
                    name=template_name
                )
                resource_url = resource.url()

                preset_resources[resource_url] = template_name

    # Sort by template name
    export_templates = dict(sorted(preset_resources.items(),
                                   key=lambda x: x[1]))

    # Add default built-ins at the start
    # TODO: find the built-ins automatically; scraped with https://gist.github.com/BigRoy/97150c7c6f0a0c916418207b9a2bc8f1  # noqa
    result = {
        "export-preset-generator://viewport2d": "2D View",  # noqa
        "export-preset-generator://doc-channel-normal-no-alpha": "Document channels + Normal + AO (No Alpha)",  # noqa
        "export-preset-generator://doc-channel-normal-with-alpha": "Document channels + Normal + AO (With Alpha)",  # noqa
        "export-preset-generator://sketchfab": "Sketchfab",  # noqa
        "export-preset-generator://adobe-standard-material": "Substance 3D Stager",  # noqa
        "export-preset-generator://usd": "USD PBR Metal Roughness",  # noqa
        "export-preset-generator://gltf": "glTF PBR Metal Roughness",  # noqa
        "export-preset-generator://gltf-displacement": "glTF PBR Metal Roughness + Displacement texture (experimental)"  # noqa
    }
    result.update(export_templates)
    return result

def get_export_templates(config, format="png", strip_folder=True):
    """Return export config outputs.

    This use the Javascript API `alg.mapexport.getPathsExportDocumentMaps`
    which returns a different output than using the Python equivalent
    `substance_painter.export.list_project_textures(config)`.

    The nice thing about the Javascript API version is that it returns the
    output textures grouped by filename template.

    A downside is that it doesn't return all the UDIM tiles but per template
    always returns a single file.

    Note:
        The file format needs to be explicitly passed to the Javascript API
        but upon exporting through the Python API the file format can be based
        on the output preset. So it's likely the file extension will mismatch

    Warning:
        Even though the function appears to solely get the expected outputs
        the Javascript API will actually create the config's texture output
        folder if it does not exist yet. As such, a valid path must be set.

    Example output:
    {
        "DefaultMaterial": {
            "$textureSet_BaseColor(_$colorSpace)(.$udim)": "DefaultMaterial_BaseColor_ACES - ACEScg.1002.png",   # noqa
            "$textureSet_Emissive(_$colorSpace)(.$udim)": "DefaultMaterial_Emissive_ACES - ACEScg.1002.png",     # noqa
            "$textureSet_Height(_$colorSpace)(.$udim)": "DefaultMaterial_Height_Utility - Raw.1002.png",         # noqa
            "$textureSet_Metallic(_$colorSpace)(.$udim)": "DefaultMaterial_Metallic_Utility - Raw.1002.png",     # noqa
            "$textureSet_Normal(_$colorSpace)(.$udim)": "DefaultMaterial_Normal_Utility - Raw.1002.png",         # noqa
            "$textureSet_Roughness(_$colorSpace)(.$udim)": "DefaultMaterial_Roughness_Utility - Raw.1002.png"    # noqa
        }
    }

    Arguments:
        config (dict) Export config
        format (str, Optional): Output format to write to, defaults to 'png'
        strip_folder (bool, Optional): Whether to strip the output folder
            from the output filenames.

    Returns:
        dict: The expected output maps.

    """
    folder = config["exportPath"].replace("\\", "/")
    preset = config["defaultExportPreset"]
    cmd = f'alg.mapexport.getPathsExportDocumentMaps("{preset}", "{folder}", "{format}")'  # noqa

    # The optional stack path argument is broken in Substance Painter 10.1
    # and fails on painter's C++ API triggering from the javascript API through
    # python. So we pass it the empty list of stack paths explicitly.
    # See `ayon-substancepainter` issue #13
    version_info = substance_painter.application.version_info()
    if version_info[0:2] >= (10, 1):
        cmd = f'alg.mapexport.getPathsExportDocumentMaps("{preset}", "{folder}", "{format}", [])'  # noqa

    result = substance_painter.js.evaluate(cmd)

    if strip_folder:
        for _stack, maps in result.items():
            for map_template, map_filepath in maps.items():
                map_filepath = map_filepath.replace("\\", "/")
                assert map_filepath.startswith(folder)
                map_filename = map_filepath[len(folder):].lstrip("/")
                maps[map_template] = map_filename

    return result

def get_filtered_export_preset(export_preset_name, channel_type_names,
                               strip_texture_set=False):
    """Return export presets included with specific channels
    requested by users.

    Args:
        export_preset_name (str): Name of export preset
        channel_type_list (list): A list of channel type requested by users
        strip_texture_set=False (bool): strip texture set name
        custom_export_preset (str): custom export preset name

    Returns:
        dict: export preset data
    """

    all_output_maps = []
    target_maps = []

    export_presets = get_export_presets()
    export_preset_nice_name = export_presets[export_preset_name]
    resource_presets = substance_painter.export.list_resource_export_presets()
    preset = next(
        (
            preset for preset in resource_presets
            if preset.resource_id.name == export_preset_nice_name
        ), None
    )
    if preset is None:
        return {}

    maps = preset.list_output_maps()
    for channel_map in maps:
        if strip_texture_set:
            old_channel_map = channel_map["fileName"]
            channel_map["fileName"] = old_channel_map.replace("_$textureSet", "")
            # export_preset_name = custom_export_preset
            all_output_maps.append(channel_map)
        else:
            all_output_maps = maps

    for channel_map in all_output_maps:
        if channel_type_names:
            for channel_name in channel_type_names:
                if not channel_map.get("fileName"):
                    continue

                if channel_name in channel_map["fileName"]:
                    target_maps.append(channel_map)
        else:
            target_maps = all_output_maps
    # Create a new preset
    return {
        "exportPresets": [
            {
                "name": export_preset_name,
                "maps": target_maps
            }
        ],
    }

def get_parsed_export_maps(config, strip_texture_set=False):
    """Return Export Config's expected output textures with parsed data.

    This tries to parse the texture outputs using a Python API export config.

    Parses template keys: $project, $mesh, $textureSet, $colorSpace, $udim

    Example:
    {("DefaultMaterial", ""): {
        "$mesh_$textureSet_BaseColor(_$colorSpace)(.$udim)": [
                {
                    // OUTPUT DATA FOR FILE #1 OF THE TEMPLATE
                },
                {
                    // OUTPUT DATA FOR FILE #2 OF THE TEMPLATE
                },
            ]
        },
    }}

    File output data (all outputs are `str`).
    1) Parsed tokens: These are parsed tokens from the template, they will
        only exist if found in the filename template and output filename.

        project: Workfile filename without extension
        mesh: Filename of the loaded mesh without extension
        textureSet: The texture set, e.g. "DefaultMaterial",
        colorSpace: The color space, e.g. "ACES - ACEScg",
        udim: The udim tile, e.g. "1001"

    2) Template output and filepath

        filepath: Full path to the resulting texture map, e.g.
            "/path/to/mesh_DefaultMaterial_BaseColor_ACES - ACEScg.1002.png",
        output: "mesh_DefaultMaterial_BaseColor_ACES - ACEScg.1002.png"
            Note: if template had slashes (folders) then `output` will too.
                  So `output` might include a folder.

    Returns:
        dict: [texture_set, stack]: {template: [file1_data, file2_data]}

    """
    # Import is here to avoid recursive lib <-> colorspace imports
    from .colorspace import get_project_channel_data

    outputs = substance_painter.export.list_project_textures(config)
    templates = get_export_templates(config, strip_folder=False)

    # Get all color spaces set for the current project
    project_colorspaces = set(
        data["colorSpace"] for data in get_project_channel_data().values()
    )

    # Get current project mesh path and project path to explicitly match
    # the $mesh and $project tokens
    project_mesh_path = substance_painter.project.last_imported_mesh_path()
    project_path = substance_painter.project.file_path() or ""

    # Get the current export path to strip this of the beginning of filepath
    # results, since filename templates don't have these we'll match without
    # that part of the filename.
    export_path = config["exportPath"]
    export_path = export_path.replace("\\", "/")
    if not export_path.endswith("/"):
        export_path += "/"

    # Parse the outputs
    result = {}
    for key, filepaths in outputs.items():
        texture_set, stack = key

        if stack:
            stack_path = f"{texture_set}/{stack}"
        else:
            stack_path = texture_set
        if strip_texture_set:
            stack_templates = list(
                template.replace("_$textureSet", "")
                for template in templates[stack_path].keys()
            )
        else:
            stack_templates = list(templates[stack_path].keys())
        template_regex = _templates_to_regex(stack_templates,
                                             texture_set=texture_set,
                                             colorspaces=project_colorspaces,
                                             mesh=project_mesh_path,
                                             project=project_path)

        # Let's precompile the regexes
        for template, regex in template_regex.items():
            template_regex[template] = re.compile(regex)

        stack_results = defaultdict(list)
        for filepath in sorted(filepaths):
            # We strip explicitly using the full parent export path instead of
            # using `os.path.basename` because export template is allowed to
            # have subfolders in its template which we want to match against
            filepath = filepath.replace("\\", "/")
            assert filepath.startswith(export_path), (
                f"Filepath {filepath} must start with folder {export_path}"
            )
            filename = filepath[len(export_path):]
            stack_results = get_stack_results(stack_results, template_regex,
                                              filename, filepath,
                                              texture_set,
                                              strip_texture_set=strip_texture_set)

        result[key] = dict(stack_results)
    if strip_texture_set:
        result = get_parsed_output_maps_as_single_output(result)

    return result

def get_parsed_output_maps_as_single_output(result):
    """Get parsed output maps as single output

    Args:
        result (dict): all parsed output maps

    Returns:
        dict: parsed output maps as single output
    """
    result_with_single_output = {}
    result_with_single_output[("", "")] = {}
    for template_maps in result.values():
        for template, outputs in template_maps.items():
            if template not in result_with_single_output[("", "")]:
                result_with_single_output[("", "")][template] = []
            result_with_single_output[("", "")][template].extend(outputs)
    return result_with_single_output

def get_stack_results(stack_results, template_regex,
                      filename, filepath,
                      texture_set,
                      strip_texture_set=False):
    """Function to get filename and filepath for parsed outputs
    """
    # Strip texture_set and stack_path if required
    if strip_texture_set:
        filename = filename.replace(f"_{texture_set}", "")
        filepath = filepath.replace(f"_{texture_set}", "")
        template_regex = {
            template.replace("_$textureSet", ""): regex
            for template, regex in template_regex.items()
        }

    # Attempt to match the filename against each template
    for template, regex in template_regex.items():
        match = regex.match(filename)
        if match:
            parsed = match.groupdict(default={})
            parsed["output"] = filename  # Add filename for convenience
            parsed["filepath"] = filepath  # Add filepath for convenience
            stack_results[template].append(parsed)
            break
    else:
        if not strip_texture_set:
            # Raise an error if no match is found
            raise ValueError(f"Unable to match {filename} against any "
                             f"template in: {list(template_regex.keys())}")
    return stack_results

def load_shelf(path, name=None):
    """Add shelf to substance painter (for current application session)

    This will dynamically add a Shelf for the current session. It's good
    to note however that these will *not* persist on restart of the host.

    Note:
        Consider the loaded shelf a static library of resources.

        The shelf will *not* be visible in application preferences in
        Edit > Settings > Libraries.

        The shelf will *not* show in the Assets browser if it has no existing
        assets

        The shelf will *not* be a selectable option for selecting it as a
        destination to import resources too.

    """

    # Ensure expanded path with forward slashes
    path = os.path.expandvars(path)
    path = os.path.abspath(path)
    path = path.replace("\\", "/")

    # Path must exist
    if not os.path.isdir(path):
        raise ValueError(f"Path is not an existing folder: {path}")

    # This name must be unique and must only contain lowercase letters,
    # numbers, underscores or hyphens.
    if name is None:
        name = os.path.basename(path)

    name = name.lower()
    name = re.sub(r"[^a-z0-9_\-]", "_", name)   # sanitize to underscores

    if substance_painter.resource.Shelves.exists(name):
        shelf = next(
            shelf for shelf in substance_painter.resource.Shelves.all()
            if shelf.name() == name
        )
        if os.path.normpath(shelf.path()) != os.path.normpath(path):
            raise ValueError(f"Shelf with name '{name}' already exists "
                             f"for a different path: '{shelf.path()}")

        return

    print(f"Adding Shelf '{name}' to path: {path}")
    substance_painter.resource.Shelves.add(name, path)

    return name

def prompt_new_file_with_mesh(mesh_filepath):
    """Prompts the user for a new file using Substance Painter's own dialog.

    This will set the mesh path to load to the given mesh and disables the
    dialog box to disallow the user to change the path. This way we can allow
    user configuration of a project but set the mesh path ourselves.

    Warning:
        This is very hacky and experimental.

    Note:
       If a project is currently open using the same mesh filepath it can't
       accurately detect whether the user had actually accepted the new project
       dialog or whether the project afterwards is still the original project,
       for example when the user might have cancelled the operation.

    """

    app = QtWidgets.QApplication.instance()
    assert os.path.isfile(mesh_filepath), \
        f"Mesh filepath does not exist: {mesh_filepath}"

    def _setup_file_dialog():
        """Set filepath in QFileDialog and trigger accept result"""
        file_dialog = app.activeModalWidget()
        assert isinstance(file_dialog, QtWidgets.QFileDialog)

        # Quickly hide the dialog
        file_dialog.hide()
        app.processEvents(QtCore.QEventLoop.ExcludeUserInputEvents, 1000)

        file_dialog.setDirectory(os.path.dirname(mesh_filepath))
        url = QtCore.QUrl.fromLocalFile(os.path.basename(mesh_filepath))
        file_dialog.selectUrl(url)
        # TODO: find a way to improve the process event to
        # load more complicated mesh
        app.processEvents(QtCore.QEventLoop.ExcludeUserInputEvents, 3000)
        file_dialog.done(file_dialog.Accepted)
        app.processEvents(QtCore.QEventLoop.AllEvents)

    def _setup_prompt():
        app.processEvents(QtCore.QEventLoop.ExcludeUserInputEvents)
        dialog = app.activeModalWidget()
        assert dialog.objectName() == "NewProjectDialog"

        # Set the window title
        mesh = os.path.basename(mesh_filepath)
        dialog.setWindowTitle(f"New Project with mesh: {mesh}")

        # Get the select mesh file button
        mesh_select = dialog.findChild(QtWidgets.QPushButton, "meshSelect")

        # Hide the select mesh button to the user to block changing of mesh
        mesh_select.setVisible(False)

        # Ensure UI is visually up-to-date
        app.processEvents(QtCore.QEventLoop.ExcludeUserInputEvents, 8000)

        # Trigger the 'select file' dialog to set the path and have the
        # new file dialog to use the path.
        QtCore.QTimer.singleShot(10, _setup_file_dialog)
        mesh_select.click()

        app.processEvents(QtCore.QEventLoop.AllEvents, 5000)

        mesh_filename = dialog.findChild(QtWidgets.QFrame, "meshFileName")
        mesh_filename_label = mesh_filename.findChild(QtWidgets.QLabel)
        if not mesh_filename_label.text():
            dialog.close()
            substance_painter.logging.warning(
                "Failed to set mesh path with the prompt dialog:"
                f"{mesh_filepath}\n\n"
                "Creating new project directly with the mesh path instead.")

    new_action = _get_new_project_action()
    if not new_action:
        raise RuntimeError("Unable to detect new file action..")

    QtCore.QTimer.singleShot(0, _setup_prompt)
    new_action.trigger()
    app.processEvents(QtCore.QEventLoop.AllEvents, 5000)

    if not substance_painter.project.is_open():
        return

    # Confirm mesh was set as expected
    project_mesh = substance_painter.project.last_imported_mesh_path()
    if os.path.normpath(project_mesh) != os.path.normpath(mesh_filepath):
        return

    return project_mesh

@contextlib.contextmanager
def set_layer_stack_opacity(node_ids, channel_types):
    """Function to set the opacity of the layer stack during
    context
    Args:
        node_ids (list[int]): Substance painter root layer node ids
        channel_types (list[str]): Channel type names as defined as
            attributes in `substance_painter.textureset.ChannelType`
    """
    # Do nothing
    if not node_ids or not channel_types:
        yield
        return

    stack = substance_painter.textureset.get_active_stack()
    stack_root_layers = (
        substance_painter.layerstack.get_root_layer_nodes(stack)
    )
    node_ids = set(node_ids)  # lookup
    excluded_nodes = [
        node for node in stack_root_layers
        if node.uid() not in node_ids
    ]

    original_opacity_values = []
    for node in excluded_nodes:
        for channel in channel_types:
            channel = channel.replace("_", "")
            chan = getattr(substance_painter.textureset.ChannelType, channel)
            original_opacity_values.append((chan, node.get_opacity(chan)))
    try:
        for node in excluded_nodes:
            for channel, _ in original_opacity_values:
                node.set_opacity(0.0, channel)
        yield
    finally:
        for node in excluded_nodes:
            for channel, opacity in original_opacity_values:
                node.set_opacity(opacity, channel)

def strip_template(template, strip="._ "):
    """Return static characters in a substance painter filename template.

    >>> strip_template("$textureSet_HELLO(.$udim)")
    # HELLO
    >>> strip_template("$mesh_$textureSet_HELLO_WORLD_$colorSpace(.$udim)")
    # HELLO_WORLD
    >>> strip_template("$textureSet_HELLO(.$udim)", strip=None)
    # _HELLO
    >>> strip_template("$mesh_$textureSet_$colorSpace(.$udim)", strip=None)
    # _HELLO_
    >>> strip_template("$textureSet_HELLO(.$udim)")
    # _HELLO

    Arguments:
        template (str): Filename template to strip.
        strip (str, optional): Characters to strip from beginning and end
            of the static string in template. Defaults to: `._ `.

    Returns:
        str: The static string in filename template.

    """
    # Return only characters that were part of the template that were static.
    # Remove all keys
    keys = ["$project", "$mesh", "$textureSet", "$udim", "$colorSpace"]
    stripped_template = template
    for key in keys:
        stripped_template = stripped_template.replace(key, "")

    # Everything inside an optional bracket space is excluded since it's not
    # static. We keep a counter to track whether we are currently iterating
    # over parts of the template that are inside an 'optional' group or not.
    counter = 0
    result = ""
    for char in stripped_template:
        if char == "(":
            counter += 1
        elif char == ")":
            counter -= 1
            if counter < 0:
                counter = 0
        else:
            if counter == 0:
                result += char

    if strip:
        # Strip of any trailing start/end characters. Technically these are
        # static but usually start and end separators like space or underscore
        # aren't wanted.
        result = result.strip(strip)

    return result