from pathlib import Path
import numpy
import bpy
import os
from typing import List, Optional
from bpy.types import Material, Operator, Context, Object, Image, Mesh, MeshUVLoopLayer, Float2AttributeValue, ShaderNodeTexImage, ShaderNodeBsdfPrincipled, ShaderNodeNormalMap
from ..core.common import SceneMatClass, MaterialListBool, ProgressTracker
from ..core.packer.rectangle_packer import MaterialImageList, BinPacker
from ..core.translations import t
from ..core.logging_setup import logger
import traceback

class MaterialImageList:
    def __init__(self):
        self.albedo: Image = None
        self.normal: Image = None
        self.emission: Image = None
        self.ambient_occlusion: Image = None
        self.height: Image = None
        self.roughness: Image = None
        self.material: Material = None
        self.parent_mesh: Object = None
        self.w: int = 0
        self.h: int = 0
        self.fit = None

def scale_images_to_largest(images: List[Image]) -> tuple[int, int]:
    x: int = 0
    y: int = 0
    
    valid_images = []
    for img in images:
        if img:
            try:
                if img.has_data:
                    valid_images.append(img)
            except ReferenceError:
                # Image has been removed from Blender's memory
                pass
    
    if not valid_images:
        return 0, 0
        
    for image in valid_images:
        x = max(x, image.size[0])
        y = max(y, image.size[1])
    
    for image in valid_images:
        image.scale(width=int(x), height=int(y))

    return x, y

def MaterialImageList_to_Image_list(classitem: MaterialImageList) -> List[Image]:
    return [
        classitem.albedo,
        classitem.normal,
        classitem.emission,
        classitem.ambient_occlusion,
        classitem.height,
        classitem.roughness
    ]

def get_material_images_from_scene(context: Context) -> list[MaterialImageList]:
    material_image_list: list[MaterialImageList] = []
    
    with ProgressTracker(context, len(context.scene.objects), "Processing Materials") as progress:
        for obj in context.scene.objects:
            if obj.type == 'MESH':
                for mat_slot in obj.material_slots:
                    # Only process materials that are selected for atlas
                    if mat_slot.material and mat_slot.material.include_in_atlas is True:
                        new_mat_image_item = MaterialImageList()
                        try:
                            new_mat_image_item.albedo = bpy.data.images[mat_slot.material.texture_atlas_albedo]
                        except Exception:
                            name = mat_slot.material.name + "_albedo_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.albedo = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.albedo.pixels[:] = numpy.tile(numpy.array([0.0,0.0,0.0,1.0]), 32*32)
                            else:
                                new_mat_image_item.albedo = bpy.data.images[name]
                        try:
                            new_mat_image_item.normal = bpy.data.images[mat_slot.material.texture_atlas_normal]
                        except Exception:
                            name = mat_slot.material.name + "_normal_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.normal = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.normal.pixels[:] = numpy.tile(numpy.array([0.5,0.5,1.0,1.0]), 32*32)
                            else:
                                new_mat_image_item.normal = bpy.data.images[name]
                        try:
                            new_mat_image_item.emission = bpy.data.images[mat_slot.material.texture_atlas_emission]
                        except Exception:
                            name = mat_slot.material.name + "_emission_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.emission = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.emission.pixels[:] = numpy.tile(numpy.array([0.0,0.0,0.0,1.0]), 32*32)
                            else:
                                new_mat_image_item.emission = bpy.data.images[name]
                        try:
                            new_mat_image_item.ambient_occlusion = bpy.data.images[mat_slot.material.texture_atlas_ambient_occlusion]
                        except Exception:
                            name = mat_slot.material.name + "_ambient_occlusion_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.ambient_occlusion = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.ambient_occlusion.pixels[:] = numpy.tile(numpy.array([1.0,1.0,1.0,1.0]), 32*32)
                            else:
                                new_mat_image_item.ambient_occlusion = bpy.data.images[name]
                        try:
                            new_mat_image_item.height = bpy.data.images[mat_slot.material.texture_atlas_height]
                        except Exception:
                            name = mat_slot.material.name + "_height_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.height = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.height.pixels[:] = numpy.tile(numpy.array([0.5,0.5,0.5,1.0]), 32*32)
                            else:
                                new_mat_image_item.height = bpy.data.images[name]
                        try:
                            new_mat_image_item.roughness = bpy.data.images[mat_slot.material.texture_atlas_roughness]
                        except Exception:
                            name = mat_slot.material.name + "_roughness_replacement"
                            if name not in bpy.data.images:
                                new_mat_image_item.roughness = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
                                new_mat_image_item.roughness.pixels[:] = numpy.tile(numpy.array([1.0,1.0,1.0,0.0]), 32*32)
                            else:
                                new_mat_image_item.roughness = bpy.data.images[name]
                        
                        new_mat_image_item.material = mat_slot.material
                        new_mat_image_item.parent_mesh = obj
                        material_image_list.append(new_mat_image_item)
                
                progress.step(f"Processed {obj.name}")
    
    return material_image_list

def prep_images_in_scene(context: Context) -> List[MaterialImageList]:
    preped_images = get_material_images_from_scene(context)
    
    with ProgressTracker(context, len(preped_images), "Preparing Images") as progress:
        for MaterialImageClass in preped_images:
            ImageList = MaterialImageList_to_Image_list(MaterialImageClass)
            MaterialImageClass.w, MaterialImageClass.h = scale_images_to_largest(ImageList)
            progress.step(f"Scaled images for {MaterialImageClass.material.name}")

    return preped_images

class AvatarToolKit_OT_AtlasMaterials(Operator):
    bl_idname = "avatar_toolkit.atlas_materials"
    bl_label = t("TextureAtlas.atlas_materials")
    bl_description = t("TextureAtlas.atlas_materials_desc")
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context: Context) -> bool:
        # Only allow operation if the file is saved and materials are selected.
        if not bpy.data.filepath:
            cls.poll_message_set(t("TextureAtlas.save_file_first"))
            return False
        return context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown
    
    def execute(self, context: Context) -> set:
        try:
            selected_materials = [m for m in prep_images_in_scene(context) 
                                if m.material and m.material.include_in_atlas]
            
            if not selected_materials:
                self.report({'WARNING'}, t("TextureAtlas.no_materials_selected"))
                return {'CANCELLED'}

            logger.info("Starting material atlas creation")
            
            packer = BinPacker(selected_materials)
            mat_images = packer.fit()

            size = [
                max([matimg.fit.w + matimg.albedo.size[0] for matimg in mat_images]),
                max([matimg.fit.h + matimg.albedo.size[1] for matimg in mat_images])
            ]
            
            atlased_mat = MaterialImageList()

            # UV Remapping
            with ProgressTracker(context, len(bpy.data.objects), "Remapping UVs") as progress:
                for mat in mat_images:
                    x, y = int(mat.fit.x), int(mat.fit.y)
                    w, h = int(mat.albedo.size[0]), int(mat.albedo.size[1])

                    for obj in bpy.data.objects:
                        if obj.type == 'MESH':
                            mesh = obj.data
                            for layer in mesh.polygons:
                                if (obj.material_slots[layer.material_index].material and 
                                    obj.material_slots[layer.material_index].material == mat.material):
                                    for loop_idx in layer.loop_indices:
                                        for layer_loops in mesh.uv_layers:
                                            uv_item = layer_loops.uv[loop_idx]
                                            uv_item.vector.x = (uv_item.vector.x*(w/size[0]))+(x/size[0])
                                            uv_item.vector.y = (uv_item.vector.y*(h/size[1]))+(y/size[1])
                        progress.step(f"Processed UVs for {obj.name}")

            # Create atlas textures
            texture_types = ["albedo", "normal", "emission", "ambient_occlusion", "height", "roughness"]
            
            with ProgressTracker(context, len(texture_types), "Creating Atlas Textures") as progress:
                for type_name in texture_types:
                    new_image_name = f"Atlas_{type_name}_{context.scene.name}_{Path(bpy.data.filepath).stem}"
                    logger.debug(f"Processing {type_name} atlas image")

                    if new_image_name in bpy.data.images:
                        bpy.data.images.remove(bpy.data.images[new_image_name])

                    canvas = bpy.data.images.new(name=new_image_name, width=int(size[0]), 
                                               height=int(size[1]), alpha=True)
                    c_w = canvas.size[0]
                    canvas_pixels = list(canvas.pixels[:])
                    
                    for mat in mat_images:
                        x, y = int(mat.fit.x), int(mat.fit.y)
                        w, h = int(mat.albedo.size[0]), int(mat.albedo.size[1])
                        image_var = getattr(mat, type_name)
                        image_pixels = list(image_var.pixels[:])
                        
                        for k in range(h):
                            for i in range(w):
                                for channel in range(4):
                                    canvas_pixels[int((((k+y)*c_w)+(i+x))*4)+channel] = \
                                        image_pixels[int(((k*w)+i)*4)+channel]

                    canvas.pixels[:] = canvas_pixels[:]
                    
                    try:
                        save_dir = os.path.dirname(bpy.data.filepath)
                        canvas.save(filepath=os.path.join(save_dir, new_image_name+".png"))
                    except Exception as save_error:
                        logger.error(f"Failed to save atlas texture: {str(save_error)}")
                        self.report({'WARNING'}, f"Could not save texture to disk, This may be due to a lack of permissions.")
                    
                    setattr(atlased_mat, type_name, canvas)
                    progress.step(f"Created {type_name} atlas")

            # Create material nodes
            atlased_mat.material = bpy.data.materials.new(
                name=f"Atlas_Final_{context.scene.name}_{Path(bpy.data.filepath).stem}")
            # Note: material.use_nodes is deprecated in Blender 5.0 - materials always use nodes
            atlased_mat.material.node_tree.nodes.clear()

            principled_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled")
            principled_node.location.x = 7.29706335067749
            principled_node.location.y = 298.918212890625

            output_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeOutputMaterial")
            output_node.location.x = 297.29705810546875
            output_node.location.y = 298.918212890625

            albedo_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            albedo_node.location.x = -588.6177978515625
            albedo_node.location.y = 414.1948547363281
            albedo_node.image = atlased_mat.albedo

            emission_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            emission_node.location.x = -588.6177978515625
            emission_node.location.y = -173.9259033203125
            emission_node.image = atlased_mat.emission

            normal_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            normal_node.location.x = -941.4189453125
            normal_node.location.y = -20.8391780853271
            normal_node.image = atlased_mat.normal

            normal_map_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeNormalMap")
            normal_map_node.location.x = -545.550537109375
            normal_map_node.location.y = -0.7543716430664062

            roughness_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            roughness_node.location.x = -592.1703491210938
            roughness_node.location.y = 206.74075317382812
            roughness_node.image = atlased_mat.roughness

            ambient_occlusion_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            ambient_occlusion_node.location.x = -906.4371337890625
            ambient_occlusion_node.location.y = -389.9602355957031
            ambient_occlusion_node.image = atlased_mat.ambient_occlusion

            height_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
            height_node.location.x = -1222.383056640625
            height_node.location.y = -375.48406982421875
            height_node.image = atlased_mat.height

            atlased_mat.material.node_tree.links.new(principled_node.inputs["Base Color"], albedo_node.outputs["Color"])
            atlased_mat.material.node_tree.links.new(principled_node.inputs["Metallic"], roughness_node.outputs["Alpha"])
            atlased_mat.material.node_tree.links.new(principled_node.inputs["Roughness"], roughness_node.outputs["Color"])
            atlased_mat.material.node_tree.links.new(principled_node.inputs["Alpha"], albedo_node.outputs["Alpha"])
            atlased_mat.material.node_tree.links.new(principled_node.inputs["Normal"], normal_map_node.outputs["Normal"])
            atlased_mat.material.node_tree.links.new(principled_node.inputs["Emission Color"], emission_node.outputs["Color"])
            atlased_mat.material.node_tree.links.new(output_node.inputs["Surface"], principled_node.outputs["BSDF"])
            atlased_mat.material.node_tree.links.new(normal_map_node.inputs["Color"], normal_node.outputs["Color"])

            # Update materials
            with ProgressTracker(context, len(context.scene.objects), "Updating Materials") as progress:
                for obj in context.scene.objects:
                    if obj.type == 'MESH':
                        mesh = obj.data
                        for i, mat_slot in enumerate(obj.material_slots):
                            if mat_slot.material and mat_slot.material.include_in_atlas:
                                mesh.materials[i] = atlased_mat.material
                        progress.step(f"Updated materials for {obj.name}")

            MaterialListBool.old_list.pop(context.scene.name, None)
            was_open = context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown
            context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown = False

            if was_open:
                bpy.ops.avatar_toolkit.expand_section_materials()
            
            for area in context.screen.areas:
                if area.type == 'VIEW_3D':
                    area.tag_redraw()

            logger.info("Material atlas creation completed successfully")
            self.report({'INFO'}, t("TextureAtlas.atlas_completed"))
            return {"FINISHED"}
            
        except Exception as e:
            logger.error(f"Error creating material atlas: {traceback.format_exc()}", exc_info=True)
            self.report({'ERROR'}, t("TextureAtlas.atlas_error"))
            raise e
