import bpy
from ..utils import showMessage
from .lightmapping_common import LightmapBakeContext, getLightmapNodeGroupName
import mathutils
import time

from ..utils_debug import debug_err, debug_log, debug_warn

class PackedObject:
    def __init__(self, object):
        self.object = object

    def beforePacking(self):
        # self.original_data = self.object.data
        # self.object.data = self.object.data.copy()

        self.___calculateSize()
        bpy.context.view_layer.objects.active = self.object
        bpy.ops.object.mode_set(mode='OBJECT')
        # create new object
        # the size is important for UV packing
        bpy.ops.mesh.primitive_plane_add(location=(0, 0, 0), size=self.size)
        self.quad = bpy.context.active_object

        self.quad.name = "__" + self.object.name + "_quad"
        self.quad.data.uv_layers[0].name = "NEEDLE_LightmapUV"
        self.quad.hide_render = True
        self.quad.hide_viewport = False # it must be active in the viewport
        for mat in self.object.data.materials:
            self.quad.data.materials.append(mat)

    def afterPacking(self):
        # self.object.data = self.original_data
        pass
        
    def ___calculateSize(self):
        vertices = self.object.data.vertices
        minVert = mathutils.Vector((9999999, 9999999, 9999999))
        maxVert = mathutils.Vector((-9999999, -9999999, -9999999))
        for loop in vertices:
            if loop.co.x < minVert.x:
                minVert.x = loop.co.x
            if loop.co.y < minVert.y:
                minVert.y = loop.co.y
            if loop.co.z < minVert.z:
                minVert.z = loop.co.z
            if loop.co.x > maxVert.x:
                maxVert.x = loop.co.x
            if loop.co.y > maxVert.y:
                maxVert.y = loop.co.y
            if loop.co.z > maxVert.z:
                maxVert.z = loop.co.z
        sizeXYZ = (maxVert - minVert)
        sizeXYZ *= self.object.scale
        self.size = sizeXYZ.x + sizeXYZ.y + sizeXYZ.z
        # self.size *= self.object.scale.x + self.object.scale.y + self.object.scale.z

        if self.object.NEEDLE_lightmapScale is not None:
            scale = self.object.NEEDLE_lightmapScale
            scale = max(0.0001, scale)
            self.size *= scale
        
        return self.size

    def select(self):
        self.object.select_set(False)
        self.quad.select_set(True)
        bpy.context.view_layer.objects.active = self.quad

    def applyUvOffsets(self):
        source_uv = self.object.data.uv_layers.get("NEEDLE_LightmapUV")
        if not source_uv:
            raise Exception("No NEEDLE_LightmapUV found on " + self.object.name)
        (srcMin, srcMax) = getUVBounds(source_uv)
        
        packedUv = self.quad.data.uv_layers[0]
        (packedMin, packedMax) = getUVBounds(packedUv)
        
        vec = mathutils.Vector((0, 0, 0, 0))
        divisor = (srcMax.x - srcMin.x)
        target_uv = self.object.data.uv_layers.get("NEEDLE_LightmapUV")
        if divisor == 0:
            # TODO: this is a hack, we should probably not pack this object at all
            target_uv.active = False
            debug_warn("Packing " + self.object.name + " failed. UV bounds are invalid, min=" + str(srcMin.x) + ", max=" + str(srcMax.x))
            return
        vec.x = (packedMax.x - packedMin.x) / divisor
        vec.y = (packedMax.y - packedMin.y) / (srcMax.y - srcMin.y)
        vec.z = packedMin.x - srcMin.x * vec.x
        vec.w = packedMin.y - srcMin.y * vec.y
        self.object["NEEDLE_lightmapOffset"] = vec

        # target_uv.active = True
        self.previous_uv_data = []
        for i in range(len(source_uv.data)):
            self.previous_uv_data.append(target_uv.data[i].uv.copy())
            target_uv.data[i].uv.x = (source_uv.data[i].uv.x) * vec.x
            target_uv.data[i].uv.y = (source_uv.data[i].uv.y) * vec.y
            target_uv.data[i].uv.x += vec.z
            target_uv.data[i].uv.y += vec.w
        pass

    def undoUVOffsets(self):
        target_uv = self.object.data.uv_layers.get("NEEDLE_LightmapUV")
        for i in range(len(target_uv.data)):
            target_uv.data[i].uv = self.previous_uv_data[i]
        pass

    def cleanup(self):
        # remove quad mesh
        # delete mesh
        # bpy.data.objects.remove(self.quad, do_unlink=True)
        bpy.data.meshes.remove(self.quad.data, do_unlink=True)
        bpy.context.view_layer.objects.active = self.object
        # self.quad.select_set(False)
        pass



class UVPacker:
    def __init__(self, objects):
        self.objects = objects
        self.packed_objects = []

    def undo(self):
        for packed_object in self.packed_objects:
            packed_object.undoUVOffsets()
        pass

    def cleanup(self):
        for packed_object in self.packed_objects:
            packed_object.cleanup()

    def pack(self, bakeContext:LightmapBakeContext):

        try:
            for object in self.objects:
                self.packed_objects.append(PackedObject(object))
            
            for obj in self.packed_objects:
                obj.beforePacking()
            
            # make sure all objects are selected
            for obj in self.packed_objects:
                obj.select()

            margin = 1

            # go into edit mode
            bpy.ops.object.mode_set(mode='EDIT')

            # # enable sync selection
            bpy.context.tool_settings.use_uv_select_sync = True

            # # select all UVs in the UV editor
            bpy.ops.mesh.select_all(action='SELECT')

            # https://docs.blender.org/api/current/bpy.ops.uv.html#bpy.ops.uv.average_islands_scale
            bpy.ops.uv.average_islands_scale(scale_uv=True)

            # # make sure there's no stretching, no rotation, and enough margin
            #if blender 3.6
            if bpy.app.version[0] >= 3 and (bpy.app.version[1] == 4 or bpy.app.version[1] == 5):
                bpy.ops.uv.lightmap_pack(PREF_CONTEXT='ALL_FACES', PREF_PACK_IN_ONE=True, PREF_NEW_UVLAYER=False, PREF_APPLY_IMAGE=False, PREF_BOX_DIV=12, PREF_MARGIN_DIV=margin)
            else:
                bpy.ops.uv.lightmap_pack(PREF_CONTEXT='ALL_FACES', PREF_PACK_IN_ONE=True, PREF_NEW_UVLAYER=False, PREF_BOX_DIV=12, PREF_MARGIN_DIV=margin)
            
            resolution = bakeContext.width
            margin = 1 / resolution * 1024 * 0.05
            print("Use lightmap island margin: " + str(margin))

            bpy.ops.uv.pack_islands(margin=margin, rotate=False)#, margin_method='FRACTION')

            # # deselect all vertices again
            bpy.ops.mesh.select_all(action='DESELECT')

            # # set to object mode
            bpy.ops.object.mode_set(mode='OBJECT')

            # apply UV offsets to objects
            for obj in self.packed_objects:
                obj.applyUvOffsets()
            
            for obj in self.packed_objects:
                obj.afterPacking()

        finally:
            self.cleanup()
            pass


def pack(objects, undo_actions:list, context:LightmapBakeContext):
    startTime = time.time()
    print("Begin UV packing")
    # get object mode
    previousMode = bpy.context.object.mode
    packer = UVPacker(objects)
    packer.pack(context)
    # reset to previous mode
    bpy.ops.object.mode_set(mode=previousMode)
    print("Finished UV packing in " + str(time.time() - startTime) + " seconds")
    if undo_actions:
        undo_actions.append({"name" : "undo lightmap uv offset", "action" : packer.undo})


def applyLightmapUVsAndRemoveMeta():
    objects_in_scene = bpy.context.scene.objects
    mode = bpy.context.mode
    if "EDIT" in mode:
        mode = "EDIT"
    bpy.ops.object.mode_set(mode='OBJECT')
    already_seen_uvs = []
    # deselect all objects
    bpy.ops.object.select_all(action='DESELECT')
    has_error = False
    for object in objects_in_scene:
        if object.type != "MESH":
            continue
        uv_layer = object.data.uv_layers.get("NEEDLE_LightmapUV")
        if uv_layer == None:
            continue
        if uv_layer in already_seen_uvs:
            showMessage(["Lightmap UVs for " + object.name + " are shared among multiple objects.", "Please make meshes unique, otherwise we can not apply the lightmap UVs to your objects.", "See Help/Needle/Logs for details"])
            debug_log("Apply lightmap UV failed because of shared UVs for " + object.name)
            object.select_set(True)
            has_error = True
            continue
        already_seen_uvs.append(uv_layer)
    
    if has_error:
        return
    for object in objects_in_scene:
        try:
            if object.type != "MESH":
                continue
            if object["NEEDLE_lightmapOffset"] == None:
                continue
            offset = object["NEEDLE_lightmapOffset"]
            uv_layer = object.data.uv_layers.get("NEEDLE_LightmapUV")
            if uv_layer == None:
                continue
            print("Applying lightmap UVs for " + object.name)
            for uv in uv_layer.data:
                ox = offset[0]
                oy = offset[1]
                oz = offset[2]
                ow = offset[3]
                uv.uv.x = uv.uv.x * ox + oz
                uv.uv.y = uv.uv.y * oy + ow
            object["NEEDLE_lightmapOffset"] = (1, 1, 0, 0)
        except:
            print("Failed to apply lightmap UVs for " + object.name + ", " + object.type)
            pass
    bpy.ops.object.mode_set(mode=mode)



def getUVBounds(uv_layer):
    # calculate bounds of UV
    min = mathutils.Vector((9999999999.0, 9999999999.0))
    max = mathutils.Vector((-9999999999.0, -9999999999.0))
    for uv in uv_layer.data:
        if uv.uv.x < min.x:
            min.x = uv.uv.x
        if uv.uv.y < min.y:
            min.y = uv.uv.y
        if uv.uv.x > max.x:
            max.x = uv.uv.x
        if uv.uv.y > max.y:
            max.y = uv.uv.y
    return (min, max)



def unwrap(object, uv_layer):
    # smart project
    prev_active = bpy.context.view_layer.objects.active
    active_uv_layer = object.data.uv_layers.active
    previous_mode = None
    if bpy.context.object:
        previous_mode = bpy.context.object.mode

    # deselect all objects
    selected_objects = [obj for obj in bpy.context.selected_objects]
    bpy.ops.object.select_all(action='DESELECT')
        
    try:
        debug_log("Begin lightmap unwrapping " + object.name)
        bpy.context.view_layer.objects.active = object
        object.data.uv_layers.active = uv_layer
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.uv.smart_project(angle_limit=1.1, island_margin=0.01, scale_to_bounds=True, correct_aspect=True, margin_method="SCALED", area_weight=1)
        debug_log("Finish lightmap unwrapping " + object.name)
    except Exception as e:
        debug_err("Failed to unwrap " + str(object.name), e)
    finally:
        bpy.ops.uv.select_all(action='DESELECT')
        # bpy.context.view_layer.objects.active = prev_active # < THIS crashes blender SOMETIMES!
        object.data.uv_layers.active = active_uv_layer
        if previous_mode is not None:
            bpy.ops.object.mode_set(mode=previous_mode)
        # restore selection
        for obj in selected_objects:
            obj.select_set(True)
        pass
