# Copyright (C) 2022 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.

import bpy, os, traceback

from bpy.types import Operator, UILayout
from bpy.props import StringProperty, BoolProperty

from typing import List, Tuple

from ..constants import ResourceType, ResourceTypes
from ..preferences import PreferencesPanel
from ..properties import Properties, PropertySection
from ..registries.resource_lists_registry import ResourceListsRegistry
from ..registries.resource_selection_registry import ResourceSelectionRegistry
from ..awp.invoke_external import render_preview, update_settings
from ..utils.previews import thumbnail_of_current_render
from ..utils.io import os_path, TempFile
from ..registries.image_packer_registry import ImagePackerRegistry
from .ui_select import ASSET_OT_ui_select

def resource_type_to_mode(rsc: ResourceType) -> str:
    if isinstance(rsc, bpy.types.Object): return 'OBJECT'
    if isinstance(rsc, bpy.types.Collection): return 'COLLECTION'
    if isinstance(rsc, bpy.types.Material): return 'MATERIAL'
    if isinstance(rsc, bpy.types.NodeTree): return 'NODE_GROUP'


class ASSET_OT_export(Operator):
    """
    Export resource.
    """
    bl_idname = 'awp.export'
    bl_label = ''
    bl_description = ''
    bl_options = {'REGISTER'}


    mode: StringProperty() # type: ignore
    filename: StringProperty() # type: ignore
    name: StringProperty() # type: ignore
    catalog: StringProperty() # type: ignore
    desc: StringProperty() # type: ignore
    author: StringProperty() # type: ignore
    tags: StringProperty() # type: ignore
    extra_tag: StringProperty() # type: ignore
    use_render_as_preview: BoolProperty() # type: ignore

    
    def __fix_volume_path(self, objects: List[bpy.types.Object]) -> List[Tuple[bpy.types.Object, str]]:
        """
        Make all paths of volumes absolute, return pairs of objects 
        and there relative paths for restoring them.
        """
        r = []
        for o in objects:
            if isinstance(o, bpy.types.Object) and o.type == 'VOLUME':
                old_volume_path = o.data.filepath
                o.data.filepath = os_path(old_volume_path)
                r.append((o, old_volume_path))

        return r


    def __export_render_preview_and_fix(self, filename: str, rsc: ResourceType, reimport_blend: str) -> bool:
        """
        Export rsc to library, render preview and apply asset settings to it.
        Returns False on failure.
        """
        export = reimport_blend if reimport_blend else filename
        mode = resource_type_to_mode(rsc)
        if not mode:
            self.report({'ERROR'}, f'Unknown mode for resource {rsc.name}')
            return False

        # Transform path to .vdb of all volumes into an absolute path, relative remap seems to fail.
        original_paths = []
        if isinstance(rsc, bpy.types.Object):
            original_paths = self.__fix_volume_path([rsc])
        elif isinstance(rsc, bpy.types.Collection):
            original_paths = self.__fix_volume_path(rsc.all_objects)

        try:
            bpy.data.libraries.write(
                export,
                set([rsc]),
                path_remap='ABSOLUTE',
                fake_user=True,
                compress=True
            )
        except Exception as ex:
            self.report({'ERROR'}, f'Failure to export Resource: {ex} (see Console)')
            print(traceback.format_exc())
            return False
        finally:
            # Restore original volume paths.
            for o, p in original_paths:
                o.data.filepath = p

        tags = self.tags.split('~')
        prefs = PreferencesPanel.get()
        with TempFile('png' if prefs.preview_image_format == 'PNG' else 'jp2') as image_file:
            preview = ''
            if self.use_render_as_preview:
                error = thumbnail_of_current_render(image_file, prefs.dimension, prefs.preview_image_format)
                if error:
                    self.report({'ERROR'}, error)
                    return False
                preview = image_file
            elif mode in [ 'MATERIAL', 'OBJECT', 'COLLECTION' ]:
                render_preview(export, image_file, mode, rsc.name)
                if os.path.exists(image_file):
                    preview = image_file
                else:
                    self.report({'ERROR'}, 'Preview rendering failed (see Console)')
                    return False
            elif mode == 'NODE_GROUP':
                if rsc.bl_idname == 'ShaderNodeTree':
                    preview = os.path.join(os.path.dirname(__file__), '..', 'data', 'images', 's-node.png')
                else:
                    preview = os.path.join(os.path.dirname(__file__), '..', 'data', 'images', 'g-node.png')
                
            update_settings(
                filename, 
                reimport_blend, 
                preview, 
                mode, 
                rsc.name, 
                self.catalog, 
                self.desc, 
                self.author, 
                tags, 
                self.extra_tag,
                ImagePackerRegistry.get().selected(),
                True
            )

            return True


    def __get_all_resources(self, context: bpy.types.Context) -> ResourceTypes:
        if self.mode == 'SELECTED_OBJECTS':
            rscs = context.selected_objects
        elif self.mode == 'RESOURCE_LIST':
            rscs = ResourceSelectionRegistry.get().selected()
        else:
            rsc = {
                'MATERIAL': bpy.data.materials,
                'OBJECT': bpy.data.objects,
                'COLLECTION': bpy.data.collections,
                'NODE_GROUP': bpy.data.node_groups,
            }[self.mode][self.name]
            rscs = [rsc]
        return rscs


    @classmethod
    def description(cls, context, properties):
        if properties.mode == 'SELECTED_OBJECTS':
            return f'Export selected Objects to "{properties.filename}"'
        elif properties.mode == 'RESOURCE_LIST':
            return f'Export selected Resources to "{properties.filename}"'
        else:
            m = {
                'MATERIAL': 'material',
                'OBJECT': 'object',
                'COLLECTION': 'collection',
                'NODE_GROUP': 'node group',
            }.get(properties.mode, '??')
            return f'Export {m} "{properties.name}" to "{properties.filename}"'


    def execute(self, context: bpy.types.Context):
        # Get ressources to export.
        rscs = self.__get_all_resources(context)
        for rsc in rscs:
            # If the file exists, we export to a temp file
            # and re-import in a background process from this temp-file.
            # Otherwise, just export and fix settings in it.
            if os.path.exists(self.filename):
                with TempFile('blend') as n:
                    result = self.__export_render_preview_and_fix(self.filename, rsc, n)
            else:
                result = self.__export_render_preview_and_fix(self.filename, rsc, '')
            
            # Break in case of export / preview gen error.
            if not result:
                break

        # Update node list.
        ResourceListsRegistry.get().update(True)

        return {'FINISHED'}


    def draw(self, context: bpy.types.Context):
        """
        Draw dialog to select which images should be packed to the library.
        """
        self.layout.ui_units_x = 40
        a = self.layout.column(align=True)
        r = a.row(align=True)
        ASSET_OT_ui_select.create_ui(r, 'CHECKBOX_HLT', 'Select All', False, 'ALL', '')
        ASSET_OT_ui_select.create_ui(r, 'CHECKBOX_DEHLT', 'Select None', False, 'NONE', '')
        g = a.grid_flow(row_major=True, columns=2, align=True)
        for k, v in ImagePackerRegistry.get().images().items():
            r = g.box().row(align=True)
            
            c = r.column(align=True)
            sr = c.row(align=True)
            ASSET_OT_ui_select.create_ui(sr, 'CHECKBOX_HLT' if v.selected else 'CHECKBOX_DEHLT', '', v.selected, 'TOGGLE', k)
            ASSET_OT_ui_select.create_ui(sr, 'QUESTION', '', v.show_info, 'INFO', k)
            if v.show_info:
                if v.image.preview and v.image.preview.icon_id:
                    c.template_icon(v.image.preview.icon_id, scale=3)
                else:
                    c.label(text='No preview')

            c = r.column(align=True)
            r = g.box().column(align=True)
            r.label(text=f'Image: {v.image.name} [ {v.image.file_format}, {v.image.generated_width}x{v.image.generated_width} ]  [ {v.image.filepath} ]')
            if v.show_info:
                for o in v.origins:
                    r.label(text=str(o))

    
    def invoke(self, context: bpy.types.Context, event):
        # Get ressources to export.
        rscs = self.__get_all_resources(context)
        if rscs:
            ImagePackerRegistry.get().update(rscs)
            if ImagePackerRegistry.get().images():
                props = Properties.get()
                # Depending on this flag, whether pack all, none or user selected images.
                if props.texture_pack_mode == 'ALL':
                    ImagePackerRegistry.get().set_all(True)
                    return self.execute(context)
                elif props.texture_pack_mode == 'NONE':
                    ImagePackerRegistry.get().set_all(False)
                    return self.execute(context)
                else:
                    return context.window_manager.invoke_props_dialog(self)
            else:
                return self.execute(context)

        return {'FINISHED'}


    @staticmethod
    def create_ui(l: UILayout, mitn: Tuple[str, str, str, str], section: PropertySection):
        if len(mitn) == 4:
            mode, icon, title, name = mitn
            extra_tag = ''
        elif len(mitn) == 5:
            mode, icon, title, name, extra_tag = mitn
        op = l.operator(ASSET_OT_export.bl_idname, text=f'Export {title}', icon=icon) # type: ASSET_OT_export
        op.mode = mode
        op.filename = section.full_filename()
        op.name = name
        op.catalog = section.export_catalog()
        op.desc = section.description
        op.author = Properties.get().author
        op.tags = section.tag_string()
        op.extra_tag = extra_tag
        op.use_render_as_preview = section.use_render_as_preview


    @staticmethod
    def call_with_resource_selection(section: PropertySection):
        """
        Invoke this operator exporting all resources in ResourceSelectionRegistry, e.g. list of collections.
        """
        bpy.ops.awp.export(
            'INVOKE_DEFAULT',
            mode='RESOURCE_LIST',
            filename=section.full_filename(),
            name='',
            catalog=section.export_catalog(),
            desc=section.description,
            author=Properties.get().author,
            tags=section.tag_string(),
            extra_tag='',
            use_render_as_preview=section.use_render_as_preview
        )
