import bpy
from .constants import image_name
import numpy as np
from .libs.PIL import Image
import mathutils

def apply_color():
    context = bpy.context
    
    paint_buffer = bpy.data.images.get(image_name)
    normal_buffer = context.scene.normal_texture_settings.texture
    custom_collection = context.scene.texture_collection_settings
    
    normal_color = context.scene.normal_color_global

    returnValue = -1
    
    if paint_buffer:
        
        # apply normal color
        if normal_buffer and normal_color != mathutils.Vector((1.0, 1.0, 1.0)):
            if paint_buffer.size[0] == normal_buffer.size[0] and paint_buffer.size[1] == normal_buffer.size[1]:
                # create arrays
                paint_mask_pixels = np.array(paint_buffer.pixels[:], dtype=np.float32).reshape(-1, 4)[:, 0]
                target_pixels = np.array(normal_buffer.pixels[:], dtype=np.float32).reshape(-1, 4)
                new_pixels = np.tile(np.array([*normal_color, 1.0], dtype=np.float32), (paint_buffer.size[0] * paint_buffer.size[1], 1))

                print("image lengths: ", len(paint_mask_pixels), len(target_pixels), len(new_pixels))
                print("image shapes: ", paint_mask_pixels.shape, target_pixels.shape, new_pixels.shape)
                print("image arrays: ", paint_mask_pixels, "next", target_pixels, "next", new_pixels)

                # Convert arrays to PIL Images
                # paint_mask_image = Image.fromarray(paint_mask_pixels[:, 3], mode='F')  # Use alpha channel for mask
                paint_mask_image = Image.fromarray(paint_mask_pixels, mode='F')  # Use alpha channel for mask
                target_image = Image.fromarray(target_pixels, mode='RGBA')
                new_image = Image.fromarray(new_pixels, mode='RGBA')

                paint_mask_image.show()
                target_image.show()
                new_image.show()

                # Blend buffers
                blended_image = Image.composite(target_image, new_image, paint_mask_image)

                # Convert blended image back to numpy array and normalize
                blended_pixels = np.asarray(blended_image).astype(np.float32)

                # Save texture
                normal_buffer.pixels[:] = blended_pixels.flatten()
                normal_buffer.update()
                
                print("Color applied")
                returnValue = 0
            
            else:
                print("Texture sizes don't match")
                returnValue = -1
        else:
            print("no normal buffer selected")
            returnValue = -1
        
        # reset paint buffer
        reset_pixels = np.array(paint_buffer.pixels[:], dtype=np.float32).reshape(-1, 4)
        reset_pixels[:] = np.array((0, 0, 0, 1), dtype=np.float32)
        paint_buffer.pixels[:] = reset_pixels.flatten()
        paint_buffer.update()
        
    else:
        print("Texture 'vPainter_buffer' not found")
        returnValue = -1
    
    return returnValue

def register():
    bpy.utils.register_class(VPAINTER_OT_apply_color)

def unregister():
    bpy.utils.unregister_class(VPAINTER_OT_apply_color)