import os
import sys

from ..blender_objects import deleteObject, deleteObjectsRecursive
from ..utils_debug import debug_err, debug_log, debug_warn
from ..utils import isMac, isWindows, showMessage
from ..blender_windows import setStatus
import bpy
from bpy_extras import view3d_utils
import mathutils
from numpy import *
import time
from .lightmapping_pack import pack, unwrap   
from .lightmapping_common import LightmapBakeContext, RenderSettings, getLightmapImage, getLightmapNodeGroupName

_lastBakeDuration = -1
def getLastBakeDuration():
    return _lastBakeDuration

def registerUndo(actions: list, name: str, undoAction: callable):
    actions.append({"name" : name, "action" : undoAction})

def isLightmapped(obj: bpy.types.Object):
    isLightmapped = getattr(obj, "NEEDLE_isLightmapped", False)
    # print(obj.name + " is lightmapped: " + str(isLightmapped))
    return isLightmapped

def getLightmapScaleOffset(obj: bpy.types.Object):
    if "NEEDLE_lightmapOffset" in obj:
        return obj["NEEDLE_lightmapOffset"]
    return None


def bakeLightmaps(settings:LightmapBakeContext, context = None):

    previousMode = None
    if bpy.context.object is not None:
        previousMode = bpy.context.object.mode
    else:
        previousMode = None
    
    # there#s also EDIT_MESH
    if previousMode is not None and "EDIT" in previousMode:
        previousMode = "EDIT"
    
    print("Begin baking lightmaps")
    setStatus("Baking lightmaps", immediate=True)
    # store selected objects
    previouslySelectedObjects = []
    for obj in bpy.context.selected_objects:
        previouslySelectedObjects.append(obj)
    activeObjects = []
    activeLights = []
    toggledLights = []
    undoActions = []
    previouslyActiveObject = bpy.context.view_layer.objects.active
    previousAreaType = bpy.context.area.type
    addedImageNodesForBaking = []
    
    renderSettings = RenderSettings() 

    # idea is to only bake selected objects and update lightmap partially
    # doesnt quite work yet tho
    onlySelected = False
    duration = 0
    preparation_duration = 0
    previousLightmap = settings.lightmap

    bpy.context.area.type = 'VIEW_3D'

    can_set_mode = bpy.context.active_object is not None

    if previousMode == "EDIT":
        # if any vertices selected:
        # if len(bpy.context.active_object.data.vertices) > 0:
        debug_warn("Can not bake in Edit Mode / this might cause problems. Deselect all...")
        # showMessage("Please switch to Object mode before baking", title="Lightmapping: Wrong Mode")
        bpy.ops.mesh.select_all(action='DESELECT')
        # bpy.ops.object.mode_set(mode='OBJECT')
        # def _delayOp():
            # window
            # op = bpy.context.window_manager.operator_properties_last("needle.bake_lightmaps")
            # op.selection_only = settings.selection_only
            # op.view_only = settings.view_only
            # op.resolution = settings.width

            # override = bpy.context.copy()
            # # override['area.type'] = 'VIEW_3D'
            # override["area"] = bpy.context.area
            # bpy.ops.needle.bake_lightmaps(override, selection_only=settings.selection_only, view_only=settings.view_only, resolution=settings.width)
            
            # bpy.ops.mesh.select_all(action='DESELECT')
            # lightmap_operator = bpy.ops.needle.bake_lightmaps
            # # lightmap_operator.selection_only = settings.selection_only
            # # lightmap_operator.view_only = settings.view_only
            # lightmap_operator('INVOKE_DEFAULT')
            # pass
        # bpy.app.timers.register(_delayOp, first_interval=0)
        # return

    # TODO: if we bake in edit mode with UV layout visible we sometimes got bugs in blender. Seems like setting object mode here is not enough
    # So let's NOT do that and ask users to enable object mode
    # if can_set_mode and bpy.context.object.mode != 'OBJECT':
    #     bpy.ops.object.mode_set(mode='OBJECT')
    # if the current mode is not object mode we cant bake
    # if can_set_mode and bpy.context.object.mode != 'OBJECT':
    #     print("WARN: Current mode is not object mode. Aborting.")
    #     showMessage("Please switch to Object mode before bak", title="Lightmapping: Wrong Mode")
    #     return

    try:
        # print(previouslyActiveObject)
        # return
        for obj in bpy.context.selected_objects:
            obj.select_set(False)
        # iterate all meshes in scene
        for obj in bpy.context.scene.objects:
            if obj.type == 'MESH':
                if isLightmapped(obj) and obj.hide_get() == False and obj.hide_viewport == False:
                    # Removed the following requirementbecause we now create a UV map once if it doesnt exist
                    # if len(obj.data.uv_layers) == 0:
                    #     print("WARN: object is missing uvs", obj.name)
                    #     continue
                    # test if the object is selected and can be baked
                    if onlySelected:
                        # check if the object is selected
                        if previouslySelectedObjects.count(obj) == 0:
                            print("Object is not selected", obj.name)
                            obj.select_set(False)
                            continue
                        elif not "NEEDLE_LightmapUV" in obj.data.uv_layers:
                            print("WARN: object is missing light map uvs")
                            obj.select_set(False)
                            continue
                    # test if the object has a material for baking
                    material = obj.active_material
                    if material is None:
                        default_material = getOrCreateDefaultLightmappingMaterial()
                        obj.active_material = default_material
                        # print("WARN: Object " + obj.name + " has no material")
                        # continue
                    # check if the object has any faces
                    if len(obj.data.polygons) == 0:
                        print("WARN: Object " + obj.name + " has no faces - will be ignored")
                        continue
                    if validateMaterialNodeTree(obj.active_material) == False:
                        print("WARN: Object " + obj.name + " has invalid material - will be ignored")
                        continue
                    obj.select_set(True)
                    activeObjects.append(obj)
                    
            if obj.type == 'LIGHT':
                # if the light is marked to be used for lightmapping we want to enable it for the baking process
                # BUT if the object is set to be invisible in the viewport we DONT want to include it in baking
                # even tho it is enabled in the component.
                isHiddenInView = obj.hide_get() or obj.hide_viewport
                if isLightmapped(obj) and isHiddenInView == False:                   
                    activeLights.append(obj)
                    # make sure it is set to be included in the render
                    if obj.hide_render == True:
                        toggledLights.append(obj)
                        obj.hide_render = False
                # if the light is to non lightmapped (or disabled in viewport) we want to disable it during baking
                elif obj.hide_render == False:
                    toggledLights.append(obj)
                    obj.hide_render = True
        
        if len(activeObjects) == 0:
            msg = "No objects in scene with lightmapping enabled and material(all lightmapped objects need a material assigned)"
            setStatus(msg, immediate=True)
            print(msg)
            showMessage(msg, title="Lightmapping: No Objects")
            debug_log(msg)
            return
        
        if len(activeLights) == 0:
            msg = "No lights in scene with lightmapping enabled"
            setStatus(msg, immediate=True)
            print(msg)
            showMessage(msg, title="Lightmapping: No Lights")
            # debug_log(msg)
            # return
        
        preparation_start_time = time.time()

        # set active object to first selected object
        bpy.context.view_layer.objects.active = activeObjects[0]

        # make sure we have a lightmap node group
        getOrCreateLightmapNodeGroup(settings)
        resetDefaultMaterial()

        if not onlySelected:
            createLightmapOffsetAndUVProperties(activeObjects, undoActions=undoActions) 
            packUVsAndCalculateLightmapOffset(activeObjects, undo_actions=undoActions, context=settings)

        # enable gpu rendering
        # TODO: register undo to set render engine back to original
        prev_render_engine = bpy.context.scene.render.engine
        prev_cycles_device = bpy.context.scene.cycles.device
        undoActions.append({"name" : "Restore render engine", "action" : lambda: setattr(bpy.context.scene.render, "engine", prev_render_engine)})
        undoActions.append({"name" : "Restore cycles device", "action" : lambda: setattr(bpy.context.scene.cycles, "device", prev_cycles_device)})
        try:
            bpy.context.scene.render.engine = 'CYCLES'
            bpy.context.scene.cycles.device = 'GPU'
            addons = bpy.context.preferences.addons
            cycles = addons['cycles']
            if isWindows():
                cycles.preferences.compute_device_type = "OPTIX"
            elif isMac():
                debug_log("MacOS: Using Metal for GPU rendering")
                cycles.preferences.compute_device_type = "METAL"
            else:
                debug_log("Linux: Using CUDA for GPU rendering")
                cycles.preferences.compute_device_type = "CUDA"
            if len(cycles.preferences.devices) > 0:
                cycles.preferences.devices[0].use = True
        except Exception as e:
            showMessage("Failed to enable GPU rendering. Make sure to select and enable a GPU device in your Blender settings", title="Lightmapping: GPU Rendering")
            debug_err(e)

        # disable denoising
        _qualityMode = getLightmapQualityMode()
        renderSettings.apply(quality_mode=_qualityMode, use_denoising=getUseDenoiser())

        # get the target image
        lightmap = settings.lightmap

        # final setup of the objects
        print("Lightmap " + str(len(activeObjects)) + " objects")
        for obj in activeObjects:
            # make sure all objects we want to bake are selected
            obj.select_set(True)
            # if the object should be used for lightmapping we need to make sure its enabled for rendering
            # otherwise we will get an error when baking
            if obj.hide_render:
                obj.hide_render = False
                registerUndo(undoActions, "Hide object " + obj.name, lambda obj=obj: setattr(obj, "hide_render", True))
            # also make sure we have a lightmap image node in each material
            # so the bake operator does use it
            for material in obj.data.materials:
                if material is None: continue
                addLightmapNodeGroupIfNotFound(material, settings)
                imageNode = material.node_tree.nodes.new(type="ShaderNodeTexImage")
                material.node_tree.nodes.active = imageNode
                imageNode.image = lightmap
                registerUndo(undoActions, "Remove image node " + imageNode.name, lambda imageNode=imageNode, material=material: material.node_tree.nodes.remove(imageNode))

        preparation_duration = round(time.time() - preparation_start_time, 2)
        debug_log("Lightmap preparation took " + str(preparation_duration) + " seconds")

        print("Begin bake...")
        setStatus("Baking bake", immediate=True)
        startTime = time.time()
        obj_count = len(activeObjects)


        # EXPERIMENTAL (VIEW ONLY / SELECTION ONLY)
        # save the previous lightmap if we want to merge it
        if settings.selection_only or settings.view_only:
            anyActive = False
            obj_count = 0
            # copy the lightmap we had before to merge it with the new result
            # previousLightmap = settings.lightmap.copy()
            for obj in activeObjects:
                if settings.view_only:
                    visibleByCamera = isObjectInCameraFrustum(obj)
                    if visibleByCamera:
                        anyActive = True
                        obj_count += 1
                        print("Visible " + obj.name)
                    obj.select_set(visibleByCamera)
                else:
                    obj.select_set(False)
            if settings.selection_only:
                for obj in previouslySelectedObjects:
                    if obj in activeObjects:
                        anyActive = True
                        obj_count += 1
                        obj.select_set(True)
                        print("Selecting " + obj.name)
            if anyActive == False:
                showMessage("No objects selected. Please select at least one object", title="Lightmapping: No Objects")
                return


        margin = 8 # settings.width / 1024 * 8
        use_clear = settings.selection_only == False and settings.view_only == False
        debug_log("Lightmapping " + str(obj_count) + " objects, margin: " + str(margin))
        bpy.ops.object.bake(type='COMBINED', 
                            uv_layer="NEEDLE_LightmapUV", 
                            margin_type="EXTEND", 
                            pass_filter={("EMIT"), ("DIRECT"), ("INDIRECT"), ("COLOR"), ("GLOSSY"), ("TRANSMISSION"),
                                        #  ("DIFFUSE") # TODO: we want this to be removed in a future update - but then our blender preview node must mix in the diffuse color
                                         },
                            margin=int(margin), 
                            use_clear=use_clear)#, save_mode='EXTERNAL', filepath="D:/image.png")
        
        duration = round(time.time() - startTime, 2)
        finishedMsg = "Finished bake in " + str(duration) + " seconds"
        print(finishedMsg)
        debug_log(finishedMsg, "Lightmap Quality Mode: " + str(_qualityMode))
        # merge lightmap with the previous image if we only bake selection
        # otherwise all other objects will suddenly be black
        # if settings.selection_only:
        #     merge_start_time = time.time()
        #     mergeLightmap(settings.lightmap, previousLightmap)
        #     merge_duration = round(time.time() - merge_start_time, 2)
        #     debug_log("Merging lightmap took " + str(merge_duration) + " seconds")
        #     duration += merge_duration

        # make sure its saved
        lightmap.pack()
        showMessage(finishedMsg, title="Lightmapping: Succeeded")
    except Exception as e:
        debug_err(e)
        showMessage(["Failed to bake lightmap", "See logs for details: Help/Needle/Open Log",  e], title="Lightmapping: Failed")
    finally:
        # restore area type
        bpy.context.area.type = previousAreaType
        # restore selection
        for obj in bpy.context.selected_objects:
            obj.select_set(False)
        
        for obj in toggledLights:
            obj.hide_render = not obj.hide_render

        # restore render settings
        renderSettings.reset()
        
        if previousLightmap is not settings.lightmap:
            debug_log("Removing temporary lightmap")
            bpy.data.images.remove(previousLightmap)
        
        # remove image nodes
        for obj, material, imageNode in addedImageNodesForBaking:
            material.node_tree.nodes.remove(imageNode)

        for undo in undoActions:
            try:
                undo["action"]()
            except Exception as e:
                print("Failed to undo action " + undo["name"])
                debug_err(e)
        
        for obj in previouslySelectedObjects:
            obj.select_set(True)

        # restore active object
        bpy.context.view_layer.objects.active = previouslyActiveObject

        settings.cleanup()

        # restore mode (it's important to do this at the end)
        # if previousMode and can_set_mode and previousMode != bpy.context.object.mode:
        #     debug_log("Restoring mode after lightmapping \"" + previousMode + "\"")
            # bpy.ops.object.mode_set(mode=previousMode)
            # bpy.app.timers.register(lambda previousMode=previousMode: __restoreMode(previousMode), first_interval=.2)

        global _lastBakeDuration
        _lastBakeDuration = duration + preparation_duration
    
    print("Finished baking lightmaps")
    setStatus("Needle Lightmapping: Finished")

def __restoreMode(mode):
    bpy.ops.object.mode_set(mode=mode)

def isObjectInCameraFrustum(obj: bpy.types.Object):
    for area in bpy.context.screen.areas:
        if area.type == 'VIEW_3D':
            for region in area.regions:
                if region.type == 'WINDOW':
                    obj_2d_coords = view3d_utils.location_3d_to_region_2d(region, area.spaces[0].region_3d, obj.location)
                    if obj_2d_coords is not None:
                        return True
    debug_log("Object " + obj.name + " is not visible by camera")
    return False

def mergeLightmap(lightmap, previousLightmap):
    # check if the size matches
    lightmap_width = lightmap.size[0]
    lightmap_height = lightmap.size[1]
    prev_width = previousLightmap.size[0]
    prev_height = previousLightmap.size[1]
    if lightmap_width != prev_width or lightmap_height != prev_height:
        debug_warn("Lightmap sizes dont match. Cannot merge lightmaps.")
        return
    # copy pixels
    lmPixels = lightmap.pixels
    previousPixels = previousLightmap.pixels
    if len(lmPixels) != len(previousPixels):
        debug_warn("Lightmap pixel count doesnt match. Cannot merge lightmaps.")
        return
    lmArr = array(lmPixels)
    previousArr = array(previousPixels)
    for i in range(0, len(lmArr), 4):
        a = lmArr[i + 3]
        if a == 0:
            lmArr[i] = previousArr[i]
            lmArr[i + 1] = previousArr[i + 1]
            lmArr[i + 2] = previousArr[i + 2]
            lmArr[i + 3] = previousArr[i + 3]

    lightmap.pixels = lmArr
    debug_log("Merged lightmap with previous lightmap")
    pass

def modifyLightmapAlpha(lightmap: bpy.types.Image, pixels):
    # lightmap.update()
    # pixels = lightmap.pixels
    # copy pixels
    from math import pi
    arr = array(pixels)
    # multiply alpha channel
    # alphaFactor = 2 # 34.49
    # arr[3::4] = arr[3::4] * alphaFactor
    for pixelIndex in range(0, len(arr), 4):
        r = arr[pixelIndex]
        g = arr[pixelIndex + 1]
        b = arr[pixelIndex + 2]
        a = arr[pixelIndex + 3]
        divisor = 8
        arr[pixelIndex] = r / divisor
        arr[pixelIndex + 1] = g / divisor
        arr[pixelIndex + 2] = b / divisor
        arr[pixelIndex + 3] = a / divisor

    lightmap.pixels = arr

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 packUVsAndCalculateLightmapOffset(objects, undo_actions, context:LightmapBakeContext):
    pack(objects, undo_actions=undo_actions, context=context)

def createLightmapOffsetAndUVProperties(objs, undoActions):
    # iterate all objects. if they have UV coordinates, make sure they have a custom property LightmapOffset that's a 4-tuple of floats
    for obj in objs:
        if obj.type == 'MESH':
            if not obj.get("NEEDLE_lightmapOffset"):
                obj["NEEDLE_lightmapOffset"] = (1.0, 1.0, 0.0, 0.0)
            # else:
            #     # TODO could also create a smart project UV here
            #     print("Object " + obj.name + " has no UV map. Aborting.")
            #     continue
            # check if object has a material. warn and abort otherwise.
            if not obj.material_slots:
                print("Object " + obj.name + " has no material. Aborting.")
                continue
    
    # iterate all objects. if they have UV map, create a new set of UV map named "NEEDLE_LightmapUV"
    # we also have to make sure there are no two objects that share the same data reference
    known_mesh_data = []
    known_lightmap_uvs = []
    for obj in objs:
        if obj.type == 'MESH':
            if obj.data.uv_layers is not None:
                
                originalData = obj.data
                has_already_seen_data = originalData in known_mesh_data

                lightmapUV = obj.data.uv_layers.get("NEEDLE_LightmapUV")
                # if this object data is shared between another object and DOES already have NEEDLE_lightmapuvs
                # we want to delete the old UVs and create new ones
                if lightmapUV and has_already_seen_data:
                    lightmapUV = None

                if not lightmapUV:
                    if len(obj.data.uv_layers) <= 1:
                        debug_log("Creating new UV map for " + obj.name)
                        lightmapUV = obj.data.uv_layers.new(name="NEEDLE_LightmapUV")
                        unwrap(obj, lightmapUV)
                    else:
                        debug_log("Object " + obj.name + " has more than one UV map. Using existing one.")
                        # temporarely rename existing UVs because we use this name in the bake operation above
                        lightmapUV = obj.data.uv_layers[1]
                        originalName = lightmapUV.name
                        lightmapUV.name = "NEEDLE_LightmapUV"
                        # if the name doesnt match we unwrap it once
                        if lightmapUV.name != originalName:
                            debug_log("Renaming UV map " + originalName + " to NEEDLE_LightmapUV → unwrap")
                            unwrap(obj, lightmapUV)
                        # if our mesh doesnt have a lot of vertices we also always unwrap:
                        # elif len(obj.data.vertices) < 100:
                            # unwrap(obj, lightmapUV)
                        else:
                            unwrap(obj, lightmapUV)
                        pass


                # we need to clone the data because multiple objects might share the same data
                clonedData = obj.data.copy()
                obj.data = clonedData

                # if id(lightmapUV.data) in known_lightmap_uvs:
                #     debug_warn("UV map \"" + lightmapUV.name + "\" is shared between multiple objects, " + str(id(lightmapUV.data)))
                #     # clone it
                #     lightmapUV.active = False
                #     obj.data.uv_layers.remove(lightmapUV)
                #     newLightmapUv = obj.data.uv_layers.new(name="NEEDLE_LightmapUV")
                #     lightmapUV = newLightmapUv
                #     unwrap(obj, lightmapUV)

                # only reset to previous data if it wasnt already in the list
                # if not has_already_seen_data:
                registerUndo(undoActions, "Undo data clone", lambda obj=obj, originalData=originalData: setattr(obj, "data", originalData))
                registerUndo(undoActions, "Undo data clone (remove clone)", lambda clonedData=clonedData: deleteObject(clonedData))
                    # pass
                
                # else:
                    # existing_layer = obj.data.uv_layers[1]
                    # lightmapUV = obj.data.uv_layers.new(name="NEEDLE_LightmapUV")
                    # # copy layer
                    # for i in range(len(existing_layer.data)):
                    #     lightmapUV.data[i].uv = existing_layer.data[i].uv
                    # # remove old layer
                    # obj.data.uv_layers.remove(existing_layer)
                    # obj.data.uv_layers.append(lightmapUV)

                # we now have lightmap uvs     
                lightmapUV.active = True
                known_lightmap_uvs.append(id(lightmapUV.data))
                known_mesh_data.append(originalData)

                # copy from existing UV map
                # for i in range(len(obj.data.uv_layers["NEEDLE_LightmapUV"].data)):
                #     obj.data.uv_layers["NEEDLE_LightmapUV"].data[i].uv = obj.data.uv_layers[0].data[i].uv

                ## this is for applying an existing calculated offset
                ## copy from existing uv map and multiply each UV coordinate by lightmapOffset
                # for i in range(len(obj.data.uv_layers["NEEDLE_LightmapUV"].data)):
                #     obj.data.uv_layers["NEEDLE_LightmapUV"].data[i].uv = obj.data.uv_layers[0].data[i].uv * vec.xy + vec.zw


# k_lightmapNodeGroupName = "NEEDLE_lightmap"

def validateMaterialNodeTree(material):
    if material is None:
        return False
    if material.node_tree is None:
        return False
    if material.node_tree.nodes is None:
        return False
    for node in material.node_tree.nodes:
        if node.label == getLightmapNodeGroupName():
            # the node tree is missing if the node was removed from bpy.data
            if node.node_tree is None:
                print("Lightmap node is missing node tree. Removing node.")
                material.node_tree.nodes.remove(node)          
    return True



__version = "NEEDLE-lightmap-nodegroup-v2"

def ensureLightmapNodeIsCorrectVersion(node: bpy.types.Node):
    """
    Checks if the lightmap node is the correct version. If not it will remove the node.
    Returns True if the node is correct, False if it was removed.
    """
    if not node:
        return False
    if not node.bl_description == __version:
        print("Lightmap node group is outdated. Removing node.")
        debug_log("Lightmap node group is outdated. Removing node.")
        bpy.data.node_groups.remove(node)
        try:
            # validate materials once to remove invalid lightmap nodes
            for material in bpy.data.materials:
                validateMaterialNodeTree(material)
        except Exception as e:
            debug_err(e)
        return False
    return True
    

def getOrCreateLightmapNodeGroup(settings:LightmapBakeContext = None):
    k_lightmapNodeGroupName = getLightmapNodeGroupName()
    lightmapNodeGroup = None

    # find image texture node
    imageTextureNode = None
    
    # check if we have a nodegroup called "Lightmap"
    for nodeGroup in bpy.data.node_groups:
        if nodeGroup.name.startswith(k_lightmapNodeGroupName):
            # destroy the node group if it's not correct
            if not ensureLightmapNodeIsCorrectVersion(nodeGroup):
                continue
            # print("Lightmap node group already exists")
            lightmapNodeGroup = nodeGroup
            for node in lightmapNodeGroup.nodes:
                if node.type == 'TEX_IMAGE':
                    imageTextureNode = node
                elif node.type == 'UVMAP':
                    node.uv_map = "NEEDLE_LightmapUV"
            break

    
    
    if imageTextureNode is not None and not settings is None:
        imageTextureNode.image = settings.lightmap
        return lightmapNodeGroup
    elif lightmapNodeGroup is not None:
        return lightmapNodeGroup

    print("Creating lightmap node group")
    
    # create a new node group called "Lightmap"
    lightmapNodeGroup = bpy.data.node_groups.new(k_lightmapNodeGroupName, "ShaderNodeTree")
    lightmapNodeGroup.name = k_lightmapNodeGroupName
    lightmapNodeGroup.bl_description = __version

    # create image texture
    imageTextureNode = lightmapNodeGroup.nodes.new("ShaderNodeTexImage")
    imageTextureNode.label = "Lightmap"
    # set lightmap image
    if settings:
        imageTextureNode.image = settings.lightmap

    # create attribte with lightmap offset
    attributeNode = lightmapNodeGroup.nodes.new("ShaderNodeAttribute")
    attributeNode.label = "NEEDLE_lightmapOffset"
    attributeNode.attribute_type = 'OBJECT'
    attributeNode.attribute_name = "NEEDLE_lightmapOffset"

    # create vector math node
    vectorMathNode = lightmapNodeGroup.nodes.new("ShaderNodeVectorMath")
    vectorMathNode.operation = 'MULTIPLY_ADD'

    # create combine XYZ
    combineXYZNode_1 = lightmapNodeGroup.nodes.new("ShaderNodeCombineXYZ")
    combineXYZNode_2 = lightmapNodeGroup.nodes.new("ShaderNodeCombineXYZ")

    # create separate XYZ
    separateXYZNode = lightmapNodeGroup.nodes.new("ShaderNodeSeparateXYZ")

    # create uvmap
    uvMapNode = lightmapNodeGroup.nodes.new("ShaderNodeUVMap")
    uvMapNode.label = "NEEDLE_LightmapUV"
    uvMapNode.uv_map = "NEEDLE_LightmapUV"

    # create mix node
    mixNode = lightmapNodeGroup.nodes.new("ShaderNodeMixRGB")
    # set color1 to white
    mixNode.inputs[1].default_value = (.5, .5, .5, 1.0)
    # create attribute with object NEEDLE_isLightmapped
    isLightmappedNode = lightmapNodeGroup.nodes.new("ShaderNodeAttribute")
    isLightmappedNode.label = "NEEDLE_isLightmapped"
    isLightmappedNode.attribute_type = 'OBJECT'
    isLightmappedNode.attribute_name = "NEEDLE_isLightmapped"
    
    # create color input node
    # inputNode = lightmapNodeGroup.nodes.new("NodeGroupInput")
    # colorInput = lightmapNodeGroup.inputs.new('NodeSocketColor','Color')
    # colorInput.default_value = (1.0, 1.0, 1.0, 1.0)
    # lightmapNodeGroup.links.new(inputNode.outputs[0], mixNode.inputs[1])
    
    # create material output
    # creating the output node inside the nodegroup so we can toggle it when we want to preview lightmaps
    # also the less we have to modify in user materials the better
    outputNode = lightmapNodeGroup.nodes.new("ShaderNodeOutputMaterial")
    previewEnabled = getLightmapPreview()
    outputNode.target = 'EEVEE' if previewEnabled else 'ALL'
    outputNode.is_active_output = True
    outputNode.label = "Output"
    outputNode.name = "Output"

    mixshader = lightmapNodeGroup.nodes.new("ShaderNodeMixShader")
    mixshader.label = "Mix Shader"
    mixshader.inputs[0].default_value = 1 / 8 # this is the factor we're using in the runtime lightmapping as well

    # multiply node connecting the isLightmapped attribute to the mix shader so we see the normal material if lightmaps are OFF for this object
    multiplyNode = lightmapNodeGroup.nodes.new("ShaderNodeMath")
    multiplyNode.operation = 'MULTIPLY'
    multiplyNode.inputs[1].default_value = mixshader.inputs[0].default_value
    lightmapNodeGroup.links.new(isLightmappedNode.outputs[0], multiplyNode.inputs[0])

    # create input shader socket
    lightmapNodeGroup.interface.new_socket(name='Shader', in_out='INPUT', socket_type='NodeSocketShader', description='Input Shader')
    inputsNode = lightmapNodeGroup.nodes.new("NodeGroupInput")
    inputShaderSocket = inputsNode.outputs[0]

    
    lightmapNodeGroup.links.new(isLightmappedNode.outputs[2], mixNode.inputs[0])
    lightmapNodeGroup.links.new(imageTextureNode.outputs[0], mixNode.inputs[2])
    # link mix shader inputs: mix from lightmap and input shader
    lightmapNodeGroup.links.new(multiplyNode.outputs[0], mixshader.inputs[0])
    lightmapNodeGroup.links.new(inputShaderSocket, mixshader.inputs[1])
    lightmapNodeGroup.links.new(mixNode.outputs[0], mixshader.inputs[2])
    # link mixshader to output
    lightmapNodeGroup.links.new(mixshader.outputs[0], outputNode.inputs[0])

    # connect multiply add to image texture vector
    lightmapNodeGroup.links.new(vectorMathNode.outputs[0], imageTextureNode.inputs[0])

    lightmapNodeGroup.links.new(uvMapNode.outputs[0], vectorMathNode.inputs[0])
    lightmapNodeGroup.links.new(combineXYZNode_1.outputs[0], vectorMathNode.inputs[1])
    lightmapNodeGroup.links.new(combineXYZNode_2.outputs[0], vectorMathNode.inputs[2])

    lightmapNodeGroup.links.new(separateXYZNode.outputs[0], combineXYZNode_1.inputs[0])
    lightmapNodeGroup.links.new(separateXYZNode.outputs[1], combineXYZNode_1.inputs[1])
    lightmapNodeGroup.links.new(separateXYZNode.outputs[2], combineXYZNode_2.inputs[0])
    # attribute to separate XYZ
    lightmapNodeGroup.links.new(attributeNode.outputs[1], separateXYZNode.inputs[0])
    # attribute alpha to combine XYZ
    lightmapNodeGroup.links.new(attributeNode.outputs[3], combineXYZNode_2.inputs[1])

    # layout nodes
    attributeNode.location = (-800, -100)
    separateXYZNode.location = (-600, 0)
    combineXYZNode_1.location = (-400, 0)
    combineXYZNode_2.location = (-400, -200)
    uvMapNode.location = (-200, 100)
    vectorMathNode.location = (0, 0)
    imageTextureNode.location = (200, 0)
    isLightmappedNode.location = (300, 200)
    inputsNode.location = (300, 350)
    mixNode.location = (600, 100)
    multiplyNode.location = (600, 400)
    mixshader.location = (800, 100)
    outputNode.location = (1000, 100)
    

k_defaultLightmappingMaterialName = "Needle Default Lightmap Material"
def getOrCreateDefaultLightmappingMaterial(allow_create: bool = True):
    materials = bpy.data.materials
    for material in materials:
        if material.name == k_defaultLightmappingMaterialName:
            addLightmapNodeGroupIfNotFound(material, None)
            return material
    if allow_create == False:
        return None
    material = materials.new(k_defaultLightmappingMaterialName)
    material.use_nodes = True
    addLightmapNodeGroupIfNotFound(material, None)
    return material

def resetDefaultMaterial():
    material = getOrCreateDefaultLightmappingMaterial(False)
    if material is None:
        return
    if material.name == k_defaultLightmappingMaterialName:
        material.use_nodes = True
        addLightmapNodeGroupIfNotFound(material, None)
        # for node in material.node_tree.nodes:
        #     node.reset()


def addLightmapNodeGroupIfNotFound(mat: bpy.types.Material, context: LightmapBakeContext):
    if mat is None:
        print("WARN: material is None")
        return

    k_lightmapNodeGroupName = getLightmapNodeGroupName()

    if mat:
        # check if material already has a lightmap node group
        for node in mat.node_tree.nodes:
            # we check via startswith because when a lightmapped object is copied, the node group name is appended with a number
            # in that case we want to remove the copied node group and insert the existing one
            # that way we dont bloat the blender file with node groups (and ideally have only ONE lightmapping node group for all materials)
            if node.type == 'GROUP' and node.node_tree and node.node_tree.name.startswith(k_lightmapNodeGroupName):
                if node.node_tree.name.startswith(k_lightmapNodeGroupName):
                    __connectMaterialOutputToNeedleLightmapShaderInputSlot(mat)
                    return
                # remove the node from the material (this node was created when the object was copied)
                mat.node_tree.nodes.remove(node)
            
    nodes = mat.node_tree.nodes
    # check if we have a nodegroup called "Lightmap"
    if k_lightmapNodeGroupName in bpy.data.node_groups:
        # find the lightmap node group
        lightmapNodeGroup = bpy.data.node_groups[k_lightmapNodeGroupName]
        # check if we have a node with the lightmap node group
        for node in nodes:
            if node.type == 'GROUP' and node.node_tree == lightmapNodeGroup:
                __connectMaterialOutputToNeedleLightmapShaderInputSlot(mat)
                return
    # create a new node group
    lightmapNodeGroup = getOrCreateLightmapNodeGroup(settings=context)
    # create a new node with the lightmap node group
    lightmapNode = nodes.new("ShaderNodeGroup")
    lightmapNode.label = k_lightmapNodeGroupName
    lightmapNode.node_tree = lightmapNodeGroup
    lightmapNode.location = (0, 0)

    # find output node
    outputNode = None
    for node in nodes:
        if node.type == 'OUTPUT_MATERIAL':
            outputNode = node
            break
    if outputNode is not None:
        # place lightmap node above output node
        lightmapNode.location = (outputNode.location[0] + 100, outputNode.location[1] + 130)

    __connectMaterialOutputToNeedleLightmapShaderInputSlot(mat)

    # make sure the lightmap node is connected to a output node
    # outputNode = nodes.new("ShaderNodeOutputMaterial")
    # outputNode.location = (lightmapNode.location[0] + 200, lightmapNode.location[1])
    # outputNode.label = "NEEDLE_output"
    # mat.node_tree.links.new(lightmapNode.outputs[0], outputNode.inputs[0])
    # # set output to EEVEE
    # outputNode.target = 'EEVEE'


def __connectMaterialOutputToNeedleLightmapShaderInputSlot(material: bpy.types.Material):
    nodes = material.node_tree.nodes
    # find output node
    outputNode = None
    for node in nodes:
        if node.type == 'OUTPUT_MATERIAL':
            outputNode = node
            break

    # find lightmap node
    lightmapNode = None
    for node in nodes:
        if node.type == 'GROUP' and node.node_tree and node.node_tree.name == getLightmapNodeGroupName():
            lightmapNode = node
            break

    # connect output node to lightmapGroupNode input
    if outputNode is not None:
        shaderInputSocket = None
        for input in lightmapNode.inputs:
            if input.type == "SHADER":
                shaderInputSocket = input
                break
        if shaderInputSocket is not None:
            outputInput = outputNode.inputs[0]
            connectedNode = None
            for link in material.node_tree.links:
                if link.to_node == outputNode and link.to_socket == outputInput:
                    connectedNode = link.from_node
                    break
            if connectedNode:
                shaderOutputSlot = None
                for output in connectedNode.outputs:
                    if output.type == "SHADER":
                        shaderOutputSlot = output
                        break
                if shaderOutputSlot:
                    material.node_tree.links.new(shaderOutputSlot, shaderInputSocket)
    pass

def onLightmapPreviewChanged(self, context):
    activeScene = bpy.context.scene
    if activeScene is None:
        return None
    val = activeScene.NEEDLE_lightmapPreview
    setLightmapPreview(val)
    return None


lightmapResolutionOptions = [
    ('128', "128", "128x128"),
    ('256', "256", "256x256"),
    ('512', "512", "512x512"),
    ('1024', "1024", "1024x1024"),
    ('2048', "2048", "2048x2048"),
    ('4096', "4096", "4096x4096"),
    ('8192', "8192", "8192x8192 (slow)"),
]
lightmapQualityMode = [
    ('PREVIEW', "Preview", "Preview: Faster baking with less quality / less samples"),
    ('HIGH', "High", "High: Slower baking with more quality / more samples"),
    ('CUSTOM', "Custom", "Custom: Use cycles settings from your scene"),
]

def register():
    bpy.types.Scene.NEEDLE_lightmapPreview = bpy.props.BoolProperty(
        name="Lightmap Preview",
        description="Toggle lightmap preview",
        default=False,
        update=onLightmapPreviewChanged
    )
    bpy.types.Scene.NEEDLE_lightmapResolution = bpy.props.EnumProperty(
        name="Lightmap Resolution",
        description="Lightmap resolution",
        items=(lightmapResolutionOptions),
        default='1024'
    )
    bpy.types.Scene.NEEDLE_lightmapQualityMode = bpy.props.EnumProperty(
        name="Lightmap Quality",
        description="Lightmap quality",
        items=(lightmapQualityMode),
        default='PREVIEW'
    )
    bpy.types.Scene.NEELDE_lightmapUseDenoiser = bpy.props.BoolProperty(
        name="Use Denoiser",
        description="Enable to use denoiser for lightmaps",
        default=False
    )


def unregister():
    del bpy.types.Scene.NEEDLE_lightmapPreview
    del bpy.types.Scene.NEEDLE_lightmapResolution
    del bpy.types.Scene.NEEDLE_lightmapQualityMode
    del bpy.types.Scene.NEELDE_lightmapUseDenoiser



def getLightmapPreview() -> bool:
    activeScene = bpy.context.scene
    if activeScene is None:
        return False
    if "NEEDLE_lightmapPreview" in activeScene:
        return activeScene["NEEDLE_lightmapPreview"]
    return False

def setLightmapPreview(preview: bool):
    activeScene = bpy.context.scene
    if activeScene is None:
        return
    activeScene["NEEDLE_lightmapPreview"] = preview

    group = getOrCreateLightmapNodeGroup()
    if group is None:
        return
    for node in group.nodes:
        if node.type == 'OUTPUT_MATERIAL':
            if preview:
                node.target = 'EEVEE'
                node.mute = False
            else:
                node.mute = True
            break

def getLightmapQualityMode() -> str:
    activeScene = bpy.context.scene
    if activeScene is None:
        return 'PREVIEW'
    if "NEEDLE_lightmapQualityMode" in activeScene:
        mode = activeScene["NEEDLE_lightmapQualityMode"]
        str = lightmapQualityMode[int(mode)][0]
        return str
    return 'PREVIEW'

def getUseDenoiser() -> bool:
    activeScene = bpy.context.scene
    if activeScene is None:
        return False
    if "NEELDE_lightmapUseDenoiser" in activeScene:
        return activeScene["NEELDE_lightmapUseDenoiser"]
    return False

def hasLightmappingEnabled(obj):
    return obj and "NEEDLE_isLightmapped" in obj and obj["NEEDLE_isLightmapped"]


def getLightmapResolutionSetting() -> int:
    activeScene = bpy.context.scene
    if activeScene is None:
        return 1024
    if "NEEDLE_lightmapResolution" in activeScene:
        res = activeScene["NEEDLE_lightmapResolution"]
        res = lightmapResolutionOptions[int(res)][0]
        return int(res)
    return 1024

def setLightmapResolutionSetting(resolution: int):
    activeScene = bpy.context.scene
    if activeScene is None:
        return
    activeScene["NEEDLE_lightmapResolution"] = resolution


    
def drawLightmapSettings(layout, context):
    activeScene = bpy.context.scene
    if activeScene is None:
        return
    isEnabled = getLightmapPreview()
    lightmapPreviewOpt = layout.row()
    isInSupportedShadingMode = True
    if getattr(context.space_data, "shading", None) and context.space_data.shading.type != "MATERIAL":
        lightmapPreviewOpt.enabled = False
        isInSupportedShadingMode = False
        if context.space_data.shading.type == "RENDERED":
            isEnabled = True
    
    label = isEnabled and "Lightmap Preview ON" or "Lightmap Preview OFF"
    lightmapPreviewOpt.prop(activeScene, "NEEDLE_lightmapPreview", text=label, toggle=True, icon='IMAGE_RGB_ALPHA')

    hint_str = "Enable preview to bake"
    if not isInSupportedShadingMode:
        hint_str = "Select Viewport: Material Preview to bake"

    if (isInSupportedShadingMode == False) and hasLightmappingEnabled(context.active_object):
        layout.label(text=hint_str, icon='INFO')
    
    opts = layout.column()
    opts.label(text="Lightmap Options")
    opts.enabled = isEnabled
    opts.prop(activeScene, "NEEDLE_lightmapResolution", text="Resolution", icon='IMAGE_RGB_ALPHA')
    opts.prop(activeScene, "NEEDLE_lightmapQualityMode", text="Quality")
    if getLightmapQualityMode() == "CUSTOM":
        cycles_settings = bpy.context.scene.cycles
        opts = opts.box()
        opts.prop(cycles_settings, "samples", text="Samples")
        opts.prop(cycles_settings, "time_limit", text="Time Limit")
        opts.prop(cycles_settings, "adaptive_threshold", text="Adaptive Threshold")
        opts.prop(cycles_settings, "use_denoising", text="Use Denoising")
        denoisingLayout = opts.column()
        denoisingLayout.enabled = cycles_settings.use_denoising
        denoisingLayout.prop(cycles_settings, "denoiser", text="Denoiser")
        denoisingLayout.prop(cycles_settings, "denoising_prefilter", text="Denoising Prefilter")
    else:
        opts.prop(activeScene, "NEELDE_lightmapUseDenoiser", text="Use Denoiser")
