# ##### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENCE BLOCK #####

bl_info = {
    "name": "Sculpt Alphas Manager",
    "description": "Displays thumbnails of alpha textures (by categories) for quicker assignment to sculpt brushes",
    "author": "Ryxx",
    "blender": (2, 81, 0),
    "version": (1, 0),
    "location": "Sculpt Mode > Properties Editor > Active Tool tab > Texture Panel",
    "category": "Sculpting"
}

import os
import sys
import bpy
import bpy.utils.previews
from bpy.types import Operator, Menu, Panel, PropertyGroup, AddonPreferences, Scene, WindowManager, BlendData
from bpy.props import StringProperty, EnumProperty
from bl_ui.properties_paint_common import brush_texture_settings
    
#--------------------------------------------------------------------------------------
# A D D O N   P R E F E R E N C E S
#--------------------------------------------------------------------------------------

class AlphasManagerPreferences(AddonPreferences):

    bl_idname = __name__

    sculpt_alphas_library: StringProperty(
        name="Sculpt Alphas",
        default="",        
        subtype='FILE_PATH',
        description = 'Main Folder containing the alphas textures used for sculpt brushes'
    )
    
    vertex_alphas_library: StringProperty(
        name="Vertex Paint Alphas",
        default="",        
        subtype='FILE_PATH',
        description = 'Main Folder containing the alphas textures used for vertex paint brushes'
    )
    
    texture_alphas_library: StringProperty(
        name="Texture Paint Alphas",
        default="",        
        subtype='FILE_PATH',
        description = 'Main Folder containing the alphas textures used for texture paint brushes'
    )        
    
    def draw(self, context):
        layout = self.layout
        col = layout.column(align=True)
        col.label(text="FOLDERS SETUP INSTRUCTIONS:")
        col.label(text="Step 1 - Create a main folder that will contain your alpha textures collections and copy paste its location in the 'Library Path' field below.")
        col.label(text="Step 2 - In that main folder, create as many sub-folders as needed and name them as you wish. These will be displayed as Categories.")
        col.label(text="Step 3 - Fill each sub-folder with you own black and white alpha textures (formats accepted: jpeg, png, tiff). These will be displayed as thumbnails.")
        col = layout.column(align=True)
        col.label(text="LIBRARY PATH:")
        col.prop(self, "sculpt_alphas_library")
        col.prop(self, "vertex_alphas_library")
        col.prop(self, "texture_alphas_library")
                
#--------------------------------------------------------------------------------------
# F U N C T I O N A L I T I E S
#--------------------------------------------------------------------------------------

# OBJECT MODE FUNCTION
def object_mode(self, context):
    object_mode = ''
    
    obj = context.object
    mode = obj.mode
    
    if mode == 'SCULPT':
        object_mode = mode
    elif mode == 'VERTEX_PAINT':
        object_mode = mode         
    elif mode == 'TEXTURE_PAINT':
        object_mode = mode
        
    return object_mode

# LIB PATH FUNCTION
def lib_path(self, context):
    lib_path = ''
    
    if object_mode(self, context) == 'SCULPT':
        lib_path = context.preferences.addons[__name__].preferences.sculpt_alphas_library
    elif object_mode(self, context) == 'VERTEX_PAINT':
        lib_path = context.preferences.addons[__name__].preferences.vertex_alphas_library         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        lib_path = context.preferences.addons[__name__].preferences.texture_alphas_library
        
    return lib_path

# LIB PATH FOLDER FUNCTION
def lib_path_folder(self, context):
    lib_path_folder = ''
    
    if object_mode(self, context) == 'SCULPT':
        lib_path_folder = 'sculpt_alphas_library'
    elif object_mode(self, context) == 'VERTEX_PAINT':
        lib_path_folder = 'vertex_alphas_library'         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        lib_path_folder = 'texture_alphas_library'
        
    return lib_path_folder

# ALPHAS CATEGORY POINTER FUNCTION
def alpha_category_pointer(self, context):
    alpha_category_pointer = ''

    if object_mode(self, context) == 'SCULPT':
        alpha_category_pointer = context.scene.category_pointer_prop.sculpt_categories
    elif object_mode(self, context) == 'VERTEX_PAINT':
        alpha_category_pointer = context.scene.category_pointer_prop.vertex_categories         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        alpha_category_pointer = context.scene.category_pointer_prop.texture_categories
        
    return alpha_category_pointer

# ALPHAS CATEGORY FUNCTION
def alpha_categories(self, context):
    alpha_categories = ''
    
    if object_mode(self, context) == 'SCULPT':
        alpha_categories = 'sculpt_categories'
    elif object_mode(self, context) == 'VERTEX_PAINT':
        alpha_categories = 'vertex_categories'         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        alpha_categories = 'texture_categories'
        
    return alpha_categories

# SELECTED ALPHA FUNCTION
def selected_alpha(self, context):
    alpha_categories = ''
    
    if object_mode(self, context) == 'SCULPT':
        selected_alpha = context.window_manager.items_in_sculpt_folder
    elif object_mode(self, context) == 'VERTEX_PAINT':
        selected_alpha = context.window_manager.items_in_vertex_folder         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        selected_alpha = context.window_manager.items_in_paint_folder
        
    return selected_alpha

# TEXTURE TYPE FUNCTION
def brush_texture_type(self, context):
    brush_texture_type = ''
    
    if object_mode(self, context) == 'SCULPT':
        brush_texture_type = context.tool_settings.sculpt.brush.texture
    elif object_mode(self, context) == 'VERTEX_PAINT':
        brush_texture_type = context.tool_settings.vertex_paint.brush.texture         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        brush_texture_type = context.tool_settings.image_paint.brush.texture
        
    return brush_texture_type

# CATEGORIES PREVIEWS FUNCTION
def preview_sub_folders_categories(self, context):
                
    list_of_category_folders = []
    no_lib_folder = [('NONE', 'None', 'None')]
    
    if not lib_path(self, context):
        return no_lib_folder

    else:        
        for folder in os.listdir(lib_path(self, context)):
            if os.path.isdir(lib_path(self, context) + folder):
                list_of_category_folders.append(folder)
        
        return [(name, name, "") for name in list_of_category_folders]
    
# CATEGORY ITEMS PREVIEWS FUNCTION
def preview_items_in_folders(self, context):
    enum_items = []

    if context is None:
        return enum_items

    wm = context.window_manager            
    directory = os.path.join(lib_path(self, context), alpha_category_pointer(self, context))

    pcoll = preview_collections["main"]

    if directory == pcoll.my_previews_dir:
        return pcoll.my_previews

    #print("Scanning directory: %s" % directory)

    if directory and os.path.exists(directory):
        image_paths = []
        for fn in os.listdir(directory):
            if fn.lower().endswith(".jpeg") or fn.lower().endswith(".jpg") or fn.lower().endswith(".png") or fn.lower().endswith(".tif"):
                image_paths.append(fn)

        for i, name in enumerate(image_paths):
            filepath = os.path.join(directory, name)
            icon = pcoll.get(name)
            if not icon:
                thumb = pcoll.load(name, filepath, 'IMAGE')
            else:
                thumb = pcoll[name]
            enum_items.append((name, name, "", thumb.icon_id, i))

    pcoll.my_previews = enum_items
    pcoll.my_previews_dir = directory
    return pcoll.my_previews

# ASSIGN BRUSH TEXTURE
def assignTexture(self, context):
    texname_no_extension = os.path.splitext(selected_alpha(self, context))[0]
    selected_alpha_path = os.path.join(lib_path(self, context), alpha_category_pointer(self, context), selected_alpha(self, context))
    
    if brush_texture_type(self, context) is not None:
        bpy.data.textures.remove(brush_texture_type(self, context), do_unlink=True, do_id_user=True, do_ui_user=True)
        
    bpy.data.images.load(selected_alpha_path, check_existing=True)
    image_to_texture = bpy.data.textures.new(texname_no_extension, 'IMAGE')
    image_to_texture.image = bpy.data.images[selected_alpha(self, context)]
    
    if object_mode(self, context) == 'SCULPT':
        bpy.context.tool_settings.sculpt.brush.texture = bpy.data.textures[texname_no_extension] 
    elif object_mode(self, context) == 'VERTEX_PAINT':
        bpy.context.tool_settings.vertex_paint.brush.texture = bpy.data.textures[texname_no_extension]         
    elif object_mode(self, context) == 'TEXTURE_PAINT':
        bpy.context.tool_settings.image_paint.brush.texture = bpy.data.textures[texname_no_extension]    

            
# CATEGORY PROPERTY SCENE
class CategoryPropertyScene(bpy.types.PropertyGroup):
    
    sculpt_categories: EnumProperty(name='Sculpt Categories', items=preview_sub_folders_categories)
    vertex_categories: EnumProperty(name='Vertex Paint Categories', items=preview_sub_folders_categories)
    texture_categories: EnumProperty(name='Texture Paint Categories', items=preview_sub_folders_categories)
            
    WindowManager.items_in_sculpt_folder = EnumProperty(items=preview_items_in_folders, update=assignTexture)
    WindowManager.items_in_vertex_folder = EnumProperty(items=preview_items_in_folders, update=assignTexture)
    WindowManager.items_in_paint_folder = EnumProperty(items=preview_items_in_folders, update=assignTexture)
        
#--------------------------------------------------------------------------------------
# T E X T U R E   P A N E L   E X T E N S I O N
#--------------------------------------------------------------------------------------
                        
# REDRAW NEW TEXTURE SETTINGS ON REGISTER           
def textureRegisterDraw(self, context):    

    showLabels = True
    iconScale = 7.5
    xScale = .15
    yScale = .2
                    
    layout = self.layout

    settings = self.paint_settings(context)
    brush = settings.brush
    mode = context.object.mode
    preferences = context.preferences.addons[__name__].preferences    
    category_pointer = context.scene.category_pointer_prop
    win = context.window_manager    
                
    col = layout.column(align=True)            
    row = col.row(align=True)

    if lib_path(self, context):    
        row.prop(category_pointer, alpha_categories(self, context), text = '')
        row.operator("open.category_folder", text = '', icon="FILE_FOLDER")    
        #row.operator("create.category_folder", text = '', icon="NEWFOLDER")
    else:
        row.prop(preferences, lib_path_folder(self, context), text = '')

    col = layout.column(align=True)
                                           
    if preview_items_in_folders(self, context) and mode == 'SCULPT':                                       
        col.template_icon_view(win, "items_in_sculpt_folder", show_labels=showLabels)
                                
    elif preview_items_in_folders(self, context) and mode == 'VERTEX_PAINT':                                      
        col.template_icon_view(win, "items_in_vertex_folder", show_labels=showLabels)
                       
    elif preview_items_in_folders(self, context) and mode == 'TEXTURE_PAINT':                     
        #split_col.operator("image.texture_refresh", text = '', icon='FILE_REFRESH')                                       
        col.template_icon_view(win, "items_in_paint_folder", show_labels=showLabels)
        
    else:
        row = col.row(align=True)                
        if lib_path(self, context):                     
            row.alignment = 'CENTER'              
            row.label(text='No textures found in ' + alpha_category_pointer(self, context) + ' category !')
            col.template_icon(icon_value=166, scale=5.0)                                    
        else:                        
            row.alignment = 'CENTER'             
            row.label(text='No library folder selected !')
            col.template_icon(icon_value=693, scale=5.0)             
                        
    #166 CHECKER                                       
    #763 IMAGE_RGB
    #2 ERROR
    #693 FILE_FOLDER
    #col = layout.column()                                            
    col.template_ID(brush, "texture", new="texture.new", live_icon=False)    

    #col = layout.column()              
    brush_texture_settings(col, brush, context.sculpt_object)


# REDRAW OLD TEXTURE SETTINGS ON UNREGISTER
def textureUnregisterDraw(self, context):    
    layout = self.layout

    settings = self.paint_settings(context)
    brush = settings.brush

    col = layout.column()

    col.template_ID_preview(brush, "texture", new="texture.new", rows=3, cols=8)

    brush_texture_settings(col, brush, context.sculpt_object)
    
#--------------------------------------------------------------------------------------
# O P E R A T O R S
#--------------------------------------------------------------------------------------

# OPEN CATEGORY FOLDER
class OpenCategoryFolder(bpy.types.Operator):
    bl_idname = "open.category_folder"
    bl_label = "Open Category Folder"
    bl_description = "Open selected category's folder in Windows explorer"
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context):              
        #os.startfile(os.path.join(lib_path(self, context), alpha_category_pointer(self, context)))
        
        if sys.platform == "win32":
            os.startfile(os.path.join(lib_path(self, context), alpha_category_pointer(self, context)))
        else:
            opener = "open" if sys.platform == "darwin" else "xdg-open"
            subprocess.call([opener, os.path.join(lib_path(self, context), alpha_category_pointer(self, context))])
                    
        return {'FINISHED'}
    
# REFRESH OPERATOR
class TextureRefresh(bpy.types.Operator):
    bl_idname = "image.texture_refresh"
    bl_label = "Texture Refresh"
    bl_description = "Refresh the textures"
        
    def execute(self, context):
                        
        for pcoll in preview_collections.values():
            bpy.utils.previews.remove(pcoll)
        preview_collections.clear()
        
        pcoll = bpy.utils.previews.new()
        pcoll.my_previews_dir = ""
        pcoll.my_previews = ()
        preview_collections["main"] = pcoll        
            
        return {'FINISHED'}
    
# CREATE NEW CATEGORY FOLDER    
class CreateCategoryFolder(bpy.types.Operator):    
    bl_idname = "create.category_folder"
    bl_label = "Create Category Folder"
    bl_description = "Create a new category's folder"    
    bl_options = {'REGISTER', 'UNDO'}

    create_category: bpy.props.StringProperty()
            
    @classmethod
    def poll(cls, context):
        return True
        
    def execute(self, context):      
                    
        createFolderPath = os.path.join(lib_path(self, context), self.create_category)

        if not os.path.exists(createFolderPath) and self.create_category is not None:
            os.makedirs(createFolderPath)
        else:
            pass

        return {'FINISHED'}
    
    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        layout = self.layout
        col = layout.column()
        col.prop(self, 'create_category', text="Name")
                                                                                                                                      
#--------------------------------------------------------------------------------------
# R E G I S T R Y
#--------------------------------------------------------------------------------------

classes = (
    AlphasManagerPreferences,
    OpenCategoryFolder,
    CategoryPropertyScene,
    TextureRefresh,
    CreateCategoryFolder,                   
)

preview_collections = {}


def register():   
    bpy.types.VIEW3D_PT_tools_brush_texture.draw = textureRegisterDraw
                
    from bpy.utils import register_class
    for cls in classes:
        register_class(cls)        
                        
    Scene.category_pointer_prop = bpy.props.PointerProperty(type=CategoryPropertyScene)    

    pcoll = bpy.utils.previews.new()
    pcoll.my_previews_dir = ""
    pcoll.my_previews = ()
    preview_collections["main"] = pcoll


def unregister():
    bpy.types.VIEW3D_PT_tools_brush_texture.draw = textureUnregisterDraw    
                          
    from bpy.utils import unregister_class
    for cls in classes:
        unregister_class(cls)
    
    del Scene.category_pointer_prop
    
    for pcoll in preview_collections.values():
        bpy.utils.previews.remove(pcoll)
    preview_collections.clear()
    
if __name__ == "__main__":
    register()