
import importlib
from pathlib import Path
import bpy
import subprocess
import gpu
from gpu_extras.batch import batch_for_shader
from mathutils import Vector
from math import sqrt

from bpy_extras.io_utils import ExportHelper
from bpy.types import Operator, PropertyGroup
from bpy.props import (BoolProperty, EnumProperty, StringProperty, IntProperty, CollectionProperty)
from .core.catalog import read_catalog
from .core.bl_utils import get_addon_prefs, unique_name, get_asset_type, get_bl_cmd, get_viewport
from .core.lib_utils import get_asset_full_path, get_asset_catalog_path, find_asset_data, clear_time_tag
from . import constants


class ASSETLIB_OT_reload_addon(Operator):
    bl_idname = "assetlibrary.reload_addon"
    bl_options = {"UNDO"}
    bl_label = 'Reload Asset Library Addon'
    bl_description = 'Reload The Asset Library Addon and the addapters'

    def execute(self, context):

        print('Execute reload', __package__)

        addon = importlib.import_module(__package__)

        addon.unregister()
        importlib.reload(addon)
        for mod in addon.modules:
            importlib.reload(mod)
        addon.register()

        return {'FINISHED'}


class ASSETLIB_OT_remove_library(Operator):
    bl_idname = "assetlibrary.remove_library"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Remove Library'
    bl_description = 'Remove Library'

    index : IntProperty(default=-1)

    def execute(self, context):
        prefs = get_addon_prefs()

        addon_lib = prefs.libraries[self.index]

        bl_libs = context.preferences.filepaths.asset_libraries
        if (bl_lib := bl_libs.get(addon_lib.name)) and bl_lib.path == addon_lib.path:
            index = list(bl_libs).index(bl_lib)
            bpy.ops.preferences.asset_library_remove(index=index)

        prefs.libraries.remove(self.index)
        return {'FINISHED'}


class ASSETLIB_OT_add_library(Operator):
    bl_idname = "assetlibrary.add_library"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Add Library'
    bl_description = 'Add Library'

    def execute(self, context):
        prefs = get_addon_prefs()

        lib = prefs.libraries.add()
        lib.expand = True
        lib.name = unique_name('Asset Library', [l.name for l in prefs.libraries])

        return {'FINISHED'}


class ASSETLIB_OT_synchronize(Operator):
    bl_idname = "assetlibrary.synchronize"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Synchronize'
    bl_description = 'Synchronize Action Lib to Local Directory'

    name : StringProperty()
    diff : StringProperty()
    blocking : BoolProperty(default=False)
    mode : EnumProperty(items=[(i.replace(' ', '_').upper(), i, '') for i in ('None', 'All', 'Auto Bundle')], default='NONE')
    directory : StringProperty(subtype='DIR_PATH')
    #conform : BoolProperty(default=False)
    #def refresh(self):
    #    for area in suitable_areas(bpy.context.screen):
    #        bpy.ops.asset.library_refresh({"area": area, 'region': area.regions[3]})
            #space_data.activate_asset_by_id(asset, deferred=deferred)

    def execute(self, context):
        prefs = get_addon_prefs()

        libs = []
        if self.name:
            libs += [prefs.libraries[self.name]]

        if self.mode == 'ALL':
            libs += prefs.libraries.values()
        elif self.mode == 'AUTO_BUNDLE':
            libs += [l for l in prefs.libraries if l.auto_bundle]
        
        if not libs:
            return {"CANCELLED"}
                
        lib_datas = [l.to_dict() for l in libs]

        print(f'Bundle Libraries: {[l.name for l in libs]}')

        script_code = dedent(f"""
            import bpy
            prefs = bpy.context.preferences.addons["asset_library"].preferences

            for lib_data in {lib_datas}:
                lib = prefs.env_libraries.add()
                lib.set_dict(lib_data)
                lib.plugin.bundle(cache_diff='{self.diff}')
            
            bpy.ops.wm.quit_blender()
            """)

        script_path = Path(bpy.app.tempdir) / 'bundle_library.py'
        script_path.write_text(script_code)

        print(script_code)

        #raise Exception()

        cmd = get_bl_cmd(script=str(script_path), background=True)

        #print(cmd) 
        if self.blocking:
            subprocess.call(cmd)
            bpy.app.timers.register(refresh_asset_browsers, first_interval=0.2)
        else:
            subprocess.Popen(cmd)
        
        return {'FINISHED'}
        

class ASSETLIB_OT_save_asset_preview(Operator):
    bl_idname = "assetlibrary.save_asset_preview"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Save Asset Preview'
    bl_description = 'Save Asset Preview'

    filepath: StringProperty(
        name="File Path",
        description="Filepath used for exporting the image",
        subtype='FILE_PATH',
    )
    check_existing: BoolProperty(
        name="Check Existing",
        description="Check and warn on overwriting existing files",
        default=True,
        options={'HIDDEN'},
    )

    quality: IntProperty(subtype='PERCENTAGE', min=0, max=100, default=90, name='Quality')

    def execute(self, context):
        prefs = get_addon_prefs()

        preview = None

        if context.asset.local_id:
            preview = context.asset.local_id.preview
            width, height = preview.image_size
            pixels = [0] * width * height * 4
            preview.image_pixels_float.foreach_get(pixels)

        else:
            asset_path = context.asset.full_library_path
            asset_type, asset_name = Path(context.asset.full_path).parts[-2:]
            asset_type = get_asset_type(asset_type)

            with bpy.data.temp_data(filepath=asset_path) as temp_data:
                with temp_data.libraries.load(asset_path, assets_only=True, link=True) as (data_from, data_to):
                    setattr(data_to, asset_type, [asset_name])
                if assets := getattr(data_to, asset_type):
                    preview = assets[0].preview
                    width, height = preview.image_size
                    # Has to read pixel in the with statement for it to work
                    pixels = [0] * width * height * 4
                    preview.image_pixels_float.foreach_get(pixels)

        if not preview:
            self.report({'ERROR'}, 'Cannot retrieve preview')
            return {"CANCELLED"}

        image = bpy.data.images.new('Asset Preview', width=width, height=height, alpha=True)
        image.pixels.foreach_set(pixels)
        try:
            image.save(filepath=self.filepath, quality=self.quality)
        except Exception as e:
            print(e)
            self.report({'ERROR'}, 'Cannot write preview')
            return {"CANCELLED"} 

        return {'FINISHED'}

    def invoke(self, context, event):
        path = Path(context.asset.name)
        if bpy.data.filepath:
            path = Path(bpy.data.filepath, context.asset.name)

        self.filepath = str(path.with_suffix('.webp'))

        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}




class ASSETLIB_OT_make_custom_preview(Operator):
    bl_idname = "assetlibrary.make_custom_preview"
    bl_label = "Custom Preview"
    bl_description = "Make a preview"

    #data_type : EnumProperty(name="Type", items=lambda s, c: constants.DATA_TYPE_ITEMS)

    def draw_border(self, context):
        if not self.is_down:
            return
            
        # 50% alpha, 2 pixel width line
        shader = gpu.shader.from_builtin('UNIFORM_COLOR')
        #gpu.state.line_width_set(1.0)
        batch = batch_for_shader(shader, 'LINE_LOOP', {"pos": self.border})
        shader.uniform_float("color", (1.0, 0.0, 0.0, 1))
        batch.draw(shader)

        # restore opengl defaults
        #gpu.state.line_width_set(1.0)
        #gpu.state.blend_set('NONE')

    def grab_view3d(self, context):
        width = int(self.release_window_pos.x - self.press_window_pos.x)
        height = width#int(self.press_window_pos.y - self.release_window_pos.y)
        x = int(self.press_window_pos.x)
        y = int(self.press_window_pos.y - width)

        print(x, y, width, height)

        scene = context.scene

        fb = gpu.state.active_framebuffer_get()
        buffer = fb.read_color(x, y, width, height, 4, 0, 'FLOAT')

        buffer.dimensions = width * height * 4

        img = bpy.data.images.get('.Asset Preview')
        if img:
            bpy.data.images.remove(img)
        img = bpy.data.images.new('.Asset Preview', width, height)
        #img.scale(width, height)
        img.pixels.foreach_set(buffer)
        img.scale(256, 256)

        pixels = [0] * 256 * 256 * 4
        img.pixels.foreach_get(pixels)

        bpy.data.images.remove(img)
        return pixels

    def modal(self, context, event):
        context.area.tag_redraw()
        
        self.mouse_pos = Vector((event.mouse_region_x, event.mouse_region_y))

        if event.type == 'LEFTMOUSE' and event.value == 'PRESS':
            self.press_window_pos = Vector((event.mouse_x, event.mouse_y))
            self.press_pos = Vector((event.mouse_region_x, event.mouse_region_y))
            print('Start Border')

            self.is_down = True

        elif event.type == 'MOUSEMOVE' and self.is_down:
            

            width = int(self.mouse_pos.x - self.press_pos.x)
            X = (self.press_pos.x-1, self.mouse_pos.x +2)
            Y = (self.press_pos.y+1, self.press_pos.y-width-2)
            #print(self.mouse_pos, self.press_pos )

            #X = sorted((self.press_pos.x, self.mouse_pos.x))
            #Y = sorted((self.press_pos.y, self.mouse_pos.y))
            #Constraint to square
            #Y[0]  = Y[1] - (X[1] - X[0])
            
            self.border = [(X[0], Y[0]), (X[1], Y[0]), (X[1], Y[1]), (X[0], Y[1])]

        elif event.type == 'LEFTMOUSE' and event.value == 'RELEASE':
            self.release_window_pos = Vector((event.mouse_x, event.mouse_y))
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, 'WINDOW')
            context.area.tag_redraw()
            self.store_preview(context)

            return {'FINISHED'}

        elif event.type in {'RIGHTMOUSE', 'ESC'}:
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}

        return {'RUNNING_MODAL'}

    def store_preview(self, context):
        asset = context.window_manager.asset_library.asset
        pixels = self.grab_view3d(context)
        asset.preview.image_size = 256, 256
        asset.preview.image_pixels_float.foreach_set(pixels)

    def invoke(self, context, event):        
        self.press_window_pos = Vector((0, 0))
        self.release_window_pos = Vector((0, 0))

        self.press_pos = Vector((0, 0))

        self.is_down = False
        self.border = []
        
        # Add the region OpenGL drawing callback
        # draw in view space with 'POST_VIEW' and 'PRE_VIEW'
        self._handle = bpy.types.SpaceView3D.draw_handler_add(self.draw_border, (context,), 'WINDOW', 'POST_PIXEL')
        
        area = get_viewport()
        region = next(r for r in area.regions if r.type =="WINDOW")
        with context.temp_override(area=area, space_data=area.spaces.active, region=region):
            context.window_manager.modal_handler_add(self)
            
        return {'RUNNING_MODAL'}
        #else:
        #    self.report({'WARNING'}, "View3D not found, cannot run operator")
        #    return {'CANCELLED'}


class ASSETLIB_OT_add_tag(Operator):
    bl_idname = "assetlibrary.tag_add"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Add Tag'
    bl_description = 'Add Tag'

    #data_type : EnumProperty(name="Type", items=lambda s, c: constants.DATA_TYPE_ITEMS)

    def execute(self, context):
        asset = context.window_manager.asset_library.asset

        new_tag = asset.asset_data.tags.new(name='Tag')
        index = list(asset.asset_data.tags).index(new_tag)
        asset.asset_data.active_tag = index

        return {"FINISHED"}


class ASSETLIB_OT_remove_tag(Operator):
    bl_idname = "assetlibrary.tag_remove"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Remove Tag'
    bl_description = 'Remove Tag'

    #data_type : EnumProperty(name="Type", items=lambda s, c: constants.DATA_TYPE_ITEMS)

    def execute(self, context):
        asset = context.window_manager.asset_library.asset
        
        if asset.asset_data.active_tag == -1:
            return {"CANCELLED"}
        
        active_tag = asset.asset_data.tags[asset.asset_data.active_tag]
        asset.asset_data.tags.remove(active_tag)
        asset.asset_data.active_tag -=1

        return {"FINISHED"}



class ASSETLIB_OT_publish_asset(Operator):
    bl_idname = "assetlibrary.publish_asset"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Publish Asset'
    bl_description = 'Publish Asset'

    name : StringProperty(name='Name')
    library : EnumProperty(name="Library", items=lambda s, c: constants.LIB_ITEMS)
    #description : StringProperty(name='Description')
    catalog : StringProperty(name='Catalog')
    data_type : EnumProperty(name="Type", items=lambda s, c: constants.DATA_TYPE_ITEMS)
    catalog_items : CollectionProperty(type=PropertyGroup)

    new_asset = False
    is_asset = False
    viewport = None
    use_overlay = False

    # @classmethod
    # def poll(self, context):
    #     return context.space_data.type == 'NODE_EDITOR' and context.space_data.edit_tree

    def invoke(self, context, event):
        if self.data_type == 'NodeTree':
            asset = context.space_data.edit_tree
        elif self.data_type == 'Material':
            asset = context.object.active_material
        elif self.data_type == 'Object':
            asset = context.object        
        
        if asset.asset_data:
            self.is_asset = True
        else:
            asset.asset_mark()
            asset.preview_ensure()
            asset.preview.image_size = 256, 256

        self.viewport = get_viewport()
        if self.viewport:
            self.use_overlay = self.viewport.spaces.active.overlay.show_overlays
            self.viewport.spaces.active.overlay.show_overlays = False

        bl_libs = context.preferences.filepaths.asset_libraries
        constants.LIB_ITEMS[:] = [(lib.name, lib.name, "") for lib in bl_libs if lib.name]

        asset_type = get_asset_type(self.data_type)
        asset_data = find_asset_data(asset.name, asset_type=asset_type, preview=True)
        
        for lib in bl_libs:
            for catalog_item in read_catalog(lib.path):
                c = self.catalog_items.add()
                c.name = catalog_item.path

        self.name = asset.name
        self.new_asset = True
        if asset_data:
            catalog = read_catalog(asset_data['library'].path)
            if catalog_item := catalog.get(id=asset_data["catalog_id"]):
                self.catalog = catalog_item.path

            self.new_asset = False
            self.library = asset_data['library'].name

            if not self.is_asset:
                if asset_data.get('preview_size'):
                    asset.preview.image_size = asset_data['preview_size']
                    asset.preview.image_pixels_float.foreach_set(asset_data['preview_pixels'])

                asset.asset_data.description = asset_data['description']

                for tag in asset_data['tags']:
                    asset.asset_data.tags.new(name=tag, skip_if_exists=True)
                clear_time_tag(asset)
        
        #asset.preview_ensure()
        context.window_manager.asset_library.asset = asset

        return context.window_manager.invoke_props_dialog(self)

    def check(self, context):
        return True

    def cancel(self, context):
        asset = context.window_manager.asset_library.asset
        if self.viewport:
            self.viewport.spaces.active.overlay.show_overlays = self.use_overlay
        if not self.is_asset:
            asset.asset_clear()

    def split_row(self, layout, name):
        split = layout.split(factor=0.225)
        split.alignment = 'RIGHT'
        split.label(text=name)
        return split

    def draw(self, context):
        asset = context.window_manager.asset_library.asset
        layout = self.layout

        col = layout.column()
        col.use_property_split = True
        col.use_property_decorate = False

        split = self.split_row(layout, "Name")
        split.prop(self, "name", text='')

        split = self.split_row(layout, "Library")
        split.prop(self, "library", text='')

        split = self.split_row(layout, "Catalog")
        split.prop_search(self, "catalog", self, "catalog_items", results_are_suggestions=True, text='')

        split = self.split_row(layout, "Description")
        split.prop(asset.asset_data, "description", text='')

        split = self.split_row(layout, "Tags")
        row = split.row()
        row.template_list("ASSETBROWSER_UL_metadata_tags", "asset_tags", asset.asset_data, "tags",
                          asset.asset_data, "active_tag", rows=3)

        col = row.column(align=True)
        col.operator("assetlibrary.tag_add", icon='ADD', text="")
        col.operator("assetlibrary.tag_remove", icon='REMOVE', text="")

        split = self.split_row(layout, "Preview")
        row = split.row()
        box = row.box()
        box.template_icon(icon_value=asset.preview.icon_id, scale=5.0)

        col = row.column(align=False)
        if self.viewport:
            col.prop(self.viewport.spaces.active.overlay, 'show_overlays', icon="OVERLAY", text="")
            col.operator("assetlibrary.make_custom_preview", icon='SCENE', text="")
            #op.data_type = self.data_type

    def execute(self, context):
        bl_libs = context.preferences.filepaths.asset_libraries
        asset = context.window_manager.asset_library.asset

        publish_library = bl_libs[self.library]
        asset_temp_blend = Path(bpy.app.tempdir, self.name).with_suffix('.blend')

        bpy.data.libraries.write(str(asset_temp_blend), {asset}, path_remap="ABSOLUTE")

        self.cancel(context) # To clear the asset mark and restore overlay

        asset_type = get_asset_type(self.data_type)
        asset_full_path = Path(asset_temp_blend, asset_type, self.name)

        cmd = get_bl_cmd(
            background=True, 
            factory_startup=True, 
            blendfile=constants.RESOURCES_DIR / 'asset_preview.blend', 
            script=constants.SCRIPTS_DIR / 'publish_library_assets.py',
            library=publish_library.path,
            assets=[asset_full_path.as_posix()],
            catalogs=[self.catalog]
            )

        print(cmd)
        subprocess.call(cmd)

        return {'FINISHED'}


class ASSETLIB_OT_publish_assets(Operator):
    bl_idname = "assetlibrary.publish_assets"
    bl_options = {"REGISTER", "UNDO"}
    bl_label = 'Publish Assets'
    bl_description = 'Publish Assets'

    library : EnumProperty(name="Library", items=lambda s, c: constants.LIB_ITEMS)
    override : BoolProperty(default=True)

    # @classmethod
    # def poll(self, context):
    #     return context.space_data.edit_tree

    def invoke(self, context, event):
        bl_libs = context.preferences.filepaths.asset_libraries
        constants.LIB_ITEMS[:] = [(lib.name, lib.name, "") for lib in bl_libs]
        
        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        layout = self.layout
        col = layout.column()
        col.use_property_split = True
        col.use_property_decorate = False

        layout.prop(self, "library")
        layout.prop(self, "override")

    def execute(self, context):
        bl_libs = context.preferences.filepaths.asset_libraries
        publish_library = bl_libs[self.library]
        preview_blend = constants.RESOURCES_DIR / 'asset_preview.blend'

        cmd = get_bl_cmd(
            background=True, 
            factory_startup=True, 
            blendfile=preview_blend, 
            script=constants.SCRIPTS_DIR / 'publish_library_assets.py',
            library=publish_library.path,
            assets=[get_asset_full_path(a) for a in context.selected_assets],
            catalogs=[get_asset_catalog_path(a) for a in context.selected_assets]
            )

        print(cmd)
        subprocess.call(cmd)
        
        return {'FINISHED'}


class ASSETLIB_OT_update_assets(Operator):
    bl_idname = 'assetlibrary.update_assets'
    bl_label = 'Update node'
    bl_options = {"REGISTER", "UNDO"}

    data_type : EnumProperty(name="Type", items=lambda s, c: constants.DATA_TYPE_ITEMS)
    selection : EnumProperty(   items=[(s, s.title(), '') for s in ('ALL', 'SELECTED', 'CURRENT')], 
                                default="CURRENT", name='All Nodes')

    @classmethod
    def poll(cls, context):
        return context.space_data.edit_tree

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        asset_libraries = context.preferences.filepaths.asset_libraries

        if self.data_type == 'NodeTree':
            assets = [context.space_data.edit_tree]
            blend_data = bpy.data.node_groups

            if self.selection == 'SELECTED':
                assets = [ n.node_tree for n in context.space_data.edit_tree.nodes 
                                    if n.type == "GROUP" and n.select]
            elif self.selection == 'ALL':
                assets =  list(bpy.data.node_groups)

        elif self.data_type == 'Material':
            asset = context.object.active_material
            blend_data = bpy.data.materials
            if self.selection == 'ALL':
                assets =  list(bpy.data.materials)

        elif self.data_type == 'Object':
            return {"CANCELLED"}

        elif self.selection == 'CURRENT':
            active_node = context.space_data.edit_tree
            assets = [active_node]
        else:
            assets = list(bpy.data.node_groups)
        
        node_names = set(n.name for n in nodes)

        for asset_library in asset_libraries:
            library_path = Path(asset_library.path)
            blend_files = [fp for fp in library_path.glob("**/*.blend") if fp.is_file()]

            node_groups = list(bpy.data.node_groups)# Storing original node_geoup to compare with imported

            link = (asset_library.import_method == 'LINK')
            for blend_file in blend_files:
                print(blend_file)
                with bpy.data.libraries.load(str(blend_file), assets_only=True, link=link) as (data_from, data_to):

                    import_node_groups = [n for n in data_from.node_groups if n in node_names]
                    print("import_node_groups", import_node_groups)
                    data_to.node_groups = import_node_groups

                    node_names -= set(import_node_groups) # Store already updated nodes

                new_node_groups = set(n for n in bpy.data.node_groups if n not in node_groups)

                for new_node_group in new_node_groups:
                    new_node_group_name = new_node_group.library_weak_reference.id_name[2:]
                    local_node_group = next((n for n in bpy.data.node_groups if n.name == new_node_group_name and n != new_node_group), None)

                    if not local_node_group:
                        print(f'No local node_group {new_node_group_name}')
                        continue
                    
                    print(f'Merge node {local_node_group.name} into {new_node_group.name}')

                    local_node_group.user_remap(new_node_group)
                    new_node_group.interface_update(context)
                    bpy.data.node_groups.remove(local_node_group)

                    new_node_group.name = new_node_group_name
                    new_node_group.asset_clear()


        return {'FINISHED'}


    def draw(self, context):
        layout = self.layout
        layout.prop(self, "selection", expand=True)

classes = (
    ASSETLIB_OT_reload_addon,
    #ASSETLIB_OT_add_library,
    #ASSETLIB_OT_remove_library,
    #ASSETLIB_OT_synchronize,
    ASSETLIB_OT_save_asset_preview,
    ASSETLIB_OT_make_custom_preview,
    ASSETLIB_OT_publish_asset,
    ASSETLIB_OT_publish_assets,
    ASSETLIB_OT_add_tag,
    ASSETLIB_OT_remove_tag
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)


def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)