use crate::ecs::generational_registry::registry_entry_by_name;
use crate::ecs::mesh::components::{
    Mesh, SkinnedVertex, Vertex, create_cone_mesh, create_cube_mesh, create_cylinder_mesh,
    create_plane_mesh, create_sphere_mesh, create_subdivided_plane_mesh, create_torus_mesh,
};
use crate::ecs::prefab::resources::mesh_cache_iter;
use crate::ecs::render_layer::components::RenderLayer;
use crate::ecs::skin::systems::{GpuSkinData, SkinningCache};
use crate::ecs::transform::queries::query_descendants;
use crate::render::wgpu::rendergraph::{PassExecutionContext, PassNode};
use std::collections::HashMap;

const INITIAL_BUFFER_SIZE: usize = 10000;
const MAX_JOINTS_PER_SKIN: usize = 256;
const MAX_SKINS: usize = 500;
const MAX_TOTAL_BONES: usize = 10_000;
const INITIAL_SKINNED_VERTEX_COUNT: usize = 100_000;
const INITIAL_SKINNED_INDEX_COUNT: usize = 300_000;

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct SelectionMaskUniforms {
    view: [[f32; 4]; 4],
    projection: [[f32; 4]; 4],
}

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct ModelMatrix {
    model: [[f32; 4]; 4],
}

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct ObjectData {
    transform_index: u32,
    mesh_id: u32,
    material_id: u32,
    batch_id: u32,
}

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct SkinnedObjectData {
    transform_index: u32,
    mesh_id: u32,
    material_id: u32,
    joint_offset: u32,
}

#[derive(Debug, Clone)]
struct SkinnedMeshData {
    index_offset: u32,
    index_count: u32,
}

pub struct SelectionMaskPass {
    pipeline: wgpu::RenderPipeline,
    uniform_buffer: wgpu::Buffer,
    uniform_bind_group: wgpu::BindGroup,
    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
    transform_buffer: wgpu::Buffer,
    object_buffer: wgpu::Buffer,
    instance_bind_group_layout: wgpu::BindGroupLayout,
    instance_bind_group: wgpu::BindGroup,
    vertex_buffer_size: usize,
    index_buffer_size: usize,
    transform_buffer_size: usize,
    object_buffer_size: usize,
    draw_calls: Vec<(u32, u32, u32)>,
    enabled: bool,
    primitives: HashMap<String, Mesh>,

    skinned_pipeline: wgpu::RenderPipeline,
    skinning_compute_pipeline: wgpu::ComputePipeline,
    skinning_compute_bind_group: wgpu::BindGroup,
    skinned_instance_bind_group_layout: wgpu::BindGroupLayout,
    skinned_instance_bind_group: wgpu::BindGroup,
    bone_transforms_buffer: wgpu::Buffer,
    inverse_bind_matrices_buffer: wgpu::Buffer,
    joint_matrices_buffer: wgpu::Buffer,
    skin_data_buffer: wgpu::Buffer,
    skinned_object_buffer: wgpu::Buffer,
    skinned_vertex_buffer: wgpu::Buffer,
    skinned_index_buffer: wgpu::Buffer,
    skinned_vertex_buffer_size: u64,
    skinned_index_buffer_size: u64,
    skinned_object_buffer_size: usize,
    skinned_meshes: HashMap<String, u32>,
    skinned_mesh_data: Vec<SkinnedMeshData>,
    current_skinned_vertex_offset: u32,
    current_skinned_index_offset: u32,
    skinned_draw_calls: Vec<(u32, u32, u32)>,
    skinning_cache: SkinningCache,
    total_joints_to_dispatch: u32,
}

impl SelectionMaskPass {
    pub fn new(device: &wgpu::Device, depth_format: wgpu::TextureFormat) -> Self {
        let shader =
            device.create_shader_module(wgpu::include_wgsl!("../../shaders/selection_mask.wgsl"));

        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Uniform Buffer"),
            size: std::mem::size_of::<SelectionMaskUniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let uniform_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Selection Mask Uniform Bind Group Layout"),
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
            });

        let uniform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Selection Mask Uniform Bind Group"),
            layout: &uniform_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: uniform_buffer.as_entire_binding(),
            }],
        });

        let vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Vertex Buffer"),
            size: (std::mem::size_of::<Vertex>() * INITIAL_BUFFER_SIZE) as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Index Buffer"),
            size: (std::mem::size_of::<u32>() * INITIAL_BUFFER_SIZE * 3) as u64,
            usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let transform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Transform Buffer"),
            size: (std::mem::size_of::<ModelMatrix>() * 100) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let object_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Object Buffer"),
            size: (std::mem::size_of::<ObjectData>() * 100) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let instance_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Selection Mask Instance Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                ],
            });

        let instance_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Selection Mask Instance Bind Group"),
            layout: &instance_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: transform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: object_buffer.as_entire_binding(),
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Selection Mask Pipeline Layout"),
            bind_group_layouts: &[&uniform_bind_group_layout, &instance_bind_group_layout],
            push_constant_ranges: &[],
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Selection Mask Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[wgpu::VertexBufferLayout {
                    array_stride: std::mem::size_of::<Vertex>() as u64,
                    step_mode: wgpu::VertexStepMode::Vertex,
                    attributes: &[
                        wgpu::VertexAttribute {
                            format: wgpu::VertexFormat::Float32x3,
                            offset: 0,
                            shader_location: 0,
                        },
                        wgpu::VertexAttribute {
                            format: wgpu::VertexFormat::Float32x3,
                            offset: 12,
                            shader_location: 1,
                        },
                        wgpu::VertexAttribute {
                            format: wgpu::VertexFormat::Float32x2,
                            offset: 24,
                            shader_location: 2,
                        },
                    ],
                }],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: wgpu::TextureFormat::R8Unorm,
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                unclipped_depth: false,
                polygon_mode: wgpu::PolygonMode::Fill,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: depth_format,
                depth_write_enabled: false,
                depth_compare: wgpu::CompareFunction::GreaterEqual,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState {
                    constant: 2,
                    slope_scale: 1.0,
                    clamp: 0.0,
                },
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        let mut primitives = HashMap::new();
        primitives.insert("Cube".to_string(), create_cube_mesh());
        primitives.insert("Sphere".to_string(), create_sphere_mesh(1.0, 16));
        primitives.insert("Plane".to_string(), create_plane_mesh(2.0));
        primitives.insert(
            "SubdividedPlane".to_string(),
            create_subdivided_plane_mesh(2.0, 20),
        );
        primitives.insert("Torus".to_string(), create_torus_mesh(1.0, 0.3, 32, 16));
        primitives.insert("Cylinder".to_string(), create_cylinder_mesh(0.5, 1.0, 16));
        primitives.insert("Cone".to_string(), create_cone_mesh(0.5, 1.0, 16));

        let skinning_compute_shader =
            device.create_shader_module(wgpu::include_wgsl!("../../shaders/skinning_compute.wgsl"));
        let skinned_mask_shader = device.create_shader_module(wgpu::include_wgsl!(
            "../../shaders/skinned_selection_mask.wgsl"
        ));

        let bone_transforms_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Bone Transforms Buffer"),
            size: (std::mem::size_of::<[[f32; 4]; 4]>() * MAX_TOTAL_BONES) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let inverse_bind_matrices_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Inverse Bind Matrices Buffer"),
            size: (std::mem::size_of::<[[f32; 4]; 4]>() * MAX_JOINTS_PER_SKIN * MAX_SKINS) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let joint_matrices_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Joint Matrices Buffer"),
            size: (std::mem::size_of::<[[f32; 4]; 4]>() * MAX_JOINTS_PER_SKIN * MAX_SKINS) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let skin_data_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Skin Data Buffer"),
            size: (std::mem::size_of::<GpuSkinData>() * MAX_SKINS) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let skinning_compute_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Selection Mask Skinning Compute Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 3,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: false },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                ],
            });

        let skinning_compute_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Selection Mask Skinning Compute Bind Group"),
            layout: &skinning_compute_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: bone_transforms_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: inverse_bind_matrices_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: skin_data_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: joint_matrices_buffer.as_entire_binding(),
                },
            ],
        });

        let skinning_compute_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Selection Mask Skinning Compute Pipeline Layout"),
                bind_group_layouts: &[&skinning_compute_bind_group_layout],
                push_constant_ranges: &[],
            });

        let skinning_compute_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("Selection Mask Skinning Compute Pipeline"),
                layout: Some(&skinning_compute_pipeline_layout),
                module: &skinning_compute_shader,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });

        let skinned_object_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Skinned Object Buffer"),
            size: (std::mem::size_of::<SkinnedObjectData>() * 100) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let skinned_instance_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Selection Mask Skinned Instance Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                ],
            });

        let skinned_instance_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Selection Mask Skinned Instance Bind Group"),
            layout: &skinned_instance_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: skinned_object_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: joint_matrices_buffer.as_entire_binding(),
                },
            ],
        });

        let skinned_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Selection Mask Skinned Pipeline Layout"),
                bind_group_layouts: &[
                    &uniform_bind_group_layout,
                    &skinned_instance_bind_group_layout,
                ],
                push_constant_ranges: &[],
            });

        let skinned_vertex_layout = wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<SkinnedVertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x3,
                },
                wgpu::VertexAttribute {
                    offset: 12,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x3,
                },
                wgpu::VertexAttribute {
                    offset: 24,
                    shader_location: 2,
                    format: wgpu::VertexFormat::Float32x2,
                },
                wgpu::VertexAttribute {
                    offset: 32,
                    shader_location: 3,
                    format: wgpu::VertexFormat::Float32x4,
                },
                wgpu::VertexAttribute {
                    offset: 48,
                    shader_location: 4,
                    format: wgpu::VertexFormat::Uint32x4,
                },
                wgpu::VertexAttribute {
                    offset: 64,
                    shader_location: 5,
                    format: wgpu::VertexFormat::Float32x4,
                },
            ],
        };

        let skinned_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Selection Mask Skinned Pipeline"),
            layout: Some(&skinned_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &skinned_mask_shader,
                entry_point: Some("vs_main"),
                buffers: &[skinned_vertex_layout],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &skinned_mask_shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: wgpu::TextureFormat::R8Unorm,
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                unclipped_depth: false,
                polygon_mode: wgpu::PolygonMode::Fill,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: depth_format,
                depth_write_enabled: false,
                depth_compare: wgpu::CompareFunction::GreaterEqual,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState {
                    constant: 2,
                    slope_scale: 1.0,
                    clamp: 0.0,
                },
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        let skinned_vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Skinned Vertex Buffer"),
            size: (std::mem::size_of::<SkinnedVertex>() * INITIAL_SKINNED_VERTEX_COUNT) as u64,
            usage: wgpu::BufferUsages::VERTEX
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let skinned_index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Selection Mask Skinned Index Buffer"),
            size: (std::mem::size_of::<u32>() * INITIAL_SKINNED_INDEX_COUNT) as u64,
            usage: wgpu::BufferUsages::INDEX
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        Self {
            pipeline,
            uniform_buffer,
            uniform_bind_group,
            vertex_buffer,
            index_buffer,
            transform_buffer,
            object_buffer,
            instance_bind_group_layout,
            instance_bind_group,
            vertex_buffer_size: INITIAL_BUFFER_SIZE,
            index_buffer_size: INITIAL_BUFFER_SIZE * 3,
            transform_buffer_size: 100,
            object_buffer_size: 100,
            draw_calls: Vec::new(),
            enabled: false,
            primitives,
            skinned_pipeline,
            skinning_compute_pipeline,
            skinning_compute_bind_group,
            skinned_instance_bind_group_layout,
            skinned_instance_bind_group,
            bone_transforms_buffer,
            inverse_bind_matrices_buffer,
            joint_matrices_buffer,
            skin_data_buffer,
            skinned_object_buffer,
            skinned_vertex_buffer,
            skinned_index_buffer,
            skinned_vertex_buffer_size: (std::mem::size_of::<SkinnedVertex>()
                * INITIAL_SKINNED_VERTEX_COUNT) as u64,
            skinned_index_buffer_size: (std::mem::size_of::<u32>() * INITIAL_SKINNED_INDEX_COUNT)
                as u64,
            skinned_object_buffer_size: 100,
            skinned_meshes: HashMap::new(),
            skinned_mesh_data: Vec::new(),
            current_skinned_vertex_offset: 0,
            current_skinned_index_offset: 0,
            skinned_draw_calls: Vec::new(),
            skinning_cache: SkinningCache::default(),
            total_joints_to_dispatch: 0,
        }
    }

    fn ensure_buffer_capacity(
        &mut self,
        device: &wgpu::Device,
        vertices: usize,
        indices: usize,
        instances: usize,
    ) {
        if vertices > self.vertex_buffer_size {
            let new_size = vertices.next_power_of_two();
            self.vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Selection Mask Vertex Buffer"),
                size: (std::mem::size_of::<Vertex>() * new_size) as u64,
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            self.vertex_buffer_size = new_size;
        }

        if indices > self.index_buffer_size {
            let new_size = indices.next_power_of_two();
            self.index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Selection Mask Index Buffer"),
                size: (std::mem::size_of::<u32>() * new_size) as u64,
                usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            self.index_buffer_size = new_size;
        }

        if instances > self.transform_buffer_size || instances > self.object_buffer_size {
            let new_size = instances.next_power_of_two().max(16);

            self.transform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Selection Mask Transform Buffer"),
                size: (std::mem::size_of::<ModelMatrix>() * new_size) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            self.transform_buffer_size = new_size;

            self.object_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Selection Mask Object Buffer"),
                size: (std::mem::size_of::<ObjectData>() * new_size) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            self.object_buffer_size = new_size;

            self.instance_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("Selection Mask Instance Bind Group"),
                layout: &self.instance_bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: self.transform_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: self.object_buffer.as_entire_binding(),
                    },
                ],
            });
        }
    }

    fn sync_skinned_meshes(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        mesh_cache: &crate::ecs::prefab::resources::MeshCache,
    ) {
        for (name, mesh) in mesh_cache_iter(mesh_cache) {
            if mesh.skin_data.is_none() {
                continue;
            }

            if self.skinned_meshes.contains_key(name.as_str()) {
                continue;
            }

            let skin_data = mesh.skin_data.as_ref().unwrap();
            let skinned_vertices = &skin_data.skinned_vertices;

            let vertex_size =
                (skinned_vertices.len() * std::mem::size_of::<SkinnedVertex>()) as u64;
            let index_size = (mesh.indices.len() * std::mem::size_of::<u32>()) as u64;

            let required_vertex_size = (self.current_skinned_vertex_offset as u64
                * std::mem::size_of::<SkinnedVertex>() as u64)
                + vertex_size;
            let required_index_size = (self.current_skinned_index_offset as u64
                * std::mem::size_of::<u32>() as u64)
                + index_size;

            if required_vertex_size > self.skinned_vertex_buffer_size {
                let new_size = (required_vertex_size * 2).max(self.skinned_vertex_buffer_size * 2);
                let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("Selection Mask Skinned Vertex Buffer (Resized)"),
                    size: new_size,
                    usage: wgpu::BufferUsages::VERTEX
                        | wgpu::BufferUsages::COPY_DST
                        | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                });

                let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("Skinned Vertex Buffer Copy"),
                });
                encoder.copy_buffer_to_buffer(
                    &self.skinned_vertex_buffer,
                    0,
                    &new_buffer,
                    0,
                    self.skinned_vertex_buffer_size,
                );
                queue.submit(Some(encoder.finish()));

                self.skinned_vertex_buffer = new_buffer;
                self.skinned_vertex_buffer_size = new_size;
            }

            if required_index_size > self.skinned_index_buffer_size {
                let new_size = (required_index_size * 2).max(self.skinned_index_buffer_size * 2);
                let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("Selection Mask Skinned Index Buffer (Resized)"),
                    size: new_size,
                    usage: wgpu::BufferUsages::INDEX
                        | wgpu::BufferUsages::COPY_DST
                        | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                });

                let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("Skinned Index Buffer Copy"),
                });
                encoder.copy_buffer_to_buffer(
                    &self.skinned_index_buffer,
                    0,
                    &new_buffer,
                    0,
                    self.skinned_index_buffer_size,
                );
                queue.submit(Some(encoder.finish()));

                self.skinned_index_buffer = new_buffer;
                self.skinned_index_buffer_size = new_size;
            }

            queue.write_buffer(
                &self.skinned_vertex_buffer,
                self.current_skinned_vertex_offset as u64
                    * std::mem::size_of::<SkinnedVertex>() as u64,
                bytemuck::cast_slice(skinned_vertices),
            );

            queue.write_buffer(
                &self.skinned_index_buffer,
                self.current_skinned_index_offset as u64 * std::mem::size_of::<u32>() as u64,
                bytemuck::cast_slice(&mesh.indices),
            );

            let mesh_id = self.skinned_mesh_data.len() as u32;
            self.skinned_mesh_data.push(SkinnedMeshData {
                index_offset: self.current_skinned_index_offset,
                index_count: mesh.indices.len() as u32,
            });

            self.skinned_meshes.insert(name.clone(), mesh_id);

            self.current_skinned_vertex_offset += skinned_vertices.len() as u32;
            self.current_skinned_index_offset += mesh.indices.len() as u32;
        }
    }

    fn rebuild_skinned_instance_bind_group(&mut self, device: &wgpu::Device) {
        self.skinned_instance_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Selection Mask Skinned Instance Bind Group"),
            layout: &self.skinned_instance_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.skinned_object_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: self.joint_matrices_buffer.as_entire_binding(),
                },
            ],
        });
    }
}

impl PassNode<crate::ecs::world::World> for SelectionMaskPass {
    fn name(&self) -> &str {
        "selection_mask_pass"
    }

    fn reads(&self) -> Vec<&str> {
        vec!["depth"]
    }

    fn writes(&self) -> Vec<&str> {
        vec!["selection_mask"]
    }

    fn prepare(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        world: &crate::ecs::world::World,
    ) {
        self.draw_calls.clear();
        self.skinned_draw_calls.clear();
        self.enabled = false;
        self.total_joints_to_dispatch = 0;

        let selected_entity = match world.resources.graphics.bounding_volume_selected_entity {
            Some(entity) => entity,
            None => return,
        };

        if !world.resources.graphics.selection_outline_enabled {
            return;
        }

        self.sync_skinned_meshes(device, queue, &world.resources.mesh_cache);

        let mut entities_to_render = vec![selected_entity];
        entities_to_render.extend(query_descendants(world, selected_entity));

        let entities_with_meshes: Vec<_> = entities_to_render
            .into_iter()
            .filter(|entity| {
                if world.get_render_mesh(*entity).is_none() {
                    return false;
                }
                if world.get_global_transform(*entity).is_none() {
                    return false;
                }
                if let Some(render_layer) = world.get_render_layer(*entity)
                    && render_layer.0 == RenderLayer::OVERLAY
                {
                    return false;
                }
                true
            })
            .collect();

        if entities_with_meshes.is_empty() {
            return;
        }

        self.enabled = true;

        let camera_matrices = match crate::ecs::camera::queries::query_active_camera_matrices(world)
        {
            Some(m) => m,
            None => return,
        };

        let uniforms = SelectionMaskUniforms {
            view: camera_matrices.view.into(),
            projection: camera_matrices.projection.into(),
        };
        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));

        let mut static_entities = Vec::new();
        let mut skinned_entities = Vec::new();

        for entity in &entities_with_meshes {
            if world.get_skin(*entity).is_some() {
                skinned_entities.push(*entity);
            } else {
                static_entities.push(*entity);
            }
        }

        let mut all_vertices: Vec<Vertex> = Vec::new();
        let mut all_indices: Vec<u32> = Vec::new();
        let mut transforms: Vec<ModelMatrix> = Vec::new();
        let mut objects: Vec<ObjectData> = Vec::new();

        let mut mesh_offsets: HashMap<String, (u32, u32, u32)> = HashMap::new();

        for entity in &static_entities {
            let render_mesh = match world.get_render_mesh(*entity) {
                Some(m) => m,
                None => continue,
            };
            let transform = match world.get_global_transform(*entity) {
                Some(t) => t,
                None => continue,
            };

            let (_vertex_offset, index_offset, index_count) =
                if let Some(&offsets) = mesh_offsets.get(&render_mesh.name) {
                    offsets
                } else {
                    let (vertices, indices) =
                        if let Some(primitive) = self.primitives.get(&render_mesh.name) {
                            (&primitive.vertices, &primitive.indices)
                        } else if let Some(mesh_entry) = registry_entry_by_name(
                            &world.resources.mesh_cache.registry,
                            &render_mesh.name,
                        ) {
                            if mesh_entry.skin_data.is_some() {
                                continue;
                            }
                            (&mesh_entry.vertices, &mesh_entry.indices)
                        } else {
                            continue;
                        };

                    let vertex_offset = all_vertices.len() as u32;
                    let index_offset = all_indices.len() as u32;

                    all_vertices.extend_from_slice(vertices);
                    all_indices.extend(indices.iter().map(|i| i + vertex_offset));

                    let index_count = indices.len() as u32;
                    mesh_offsets.insert(
                        render_mesh.name.clone(),
                        (vertex_offset, index_offset, index_count),
                    );
                    (vertex_offset, index_offset, index_count)
                };

            let transform_index = transforms.len() as u32;
            let model_matrix: [[f32; 4]; 4] = transform.0.into();

            transforms.push(ModelMatrix {
                model: model_matrix,
            });

            objects.push(ObjectData {
                transform_index,
                mesh_id: 0,
                material_id: 0,
                batch_id: 0,
            });

            self.draw_calls
                .push((index_count, index_offset, transforms.len() as u32 - 1));
        }

        if !all_vertices.is_empty() && !all_indices.is_empty() {
            self.ensure_buffer_capacity(
                device,
                all_vertices.len(),
                all_indices.len(),
                transforms.len(),
            );

            queue.write_buffer(&self.vertex_buffer, 0, bytemuck::cast_slice(&all_vertices));
            queue.write_buffer(&self.index_buffer, 0, bytemuck::cast_slice(&all_indices));
            queue.write_buffer(&self.transform_buffer, 0, bytemuck::cast_slice(&transforms));
            queue.write_buffer(&self.object_buffer, 0, bytemuck::cast_slice(&objects));
        }

        if !skinned_entities.is_empty() {
            if self.skinning_cache.needs_rebuild(world) {
                self.skinning_cache.rebuild_static_data(world);

                if !self.skinning_cache.inverse_bind_matrices.is_empty() {
                    queue.write_buffer(
                        &self.inverse_bind_matrices_buffer,
                        0,
                        bytemuck::cast_slice(&self.skinning_cache.inverse_bind_matrices),
                    );
                }

                if !self.skinning_cache.skin_data.is_empty() {
                    queue.write_buffer(
                        &self.skin_data_buffer,
                        0,
                        bytemuck::cast_slice(&self.skinning_cache.skin_data),
                    );
                }
            }

            let bone_transforms = self.skinning_cache.collect_bone_transforms(world);
            if !bone_transforms.is_empty() {
                queue.write_buffer(
                    &self.bone_transforms_buffer,
                    0,
                    bytemuck::cast_slice(&bone_transforms),
                );
            }

            self.total_joints_to_dispatch = self.skinning_cache.total_joints;

            let mut skinned_objects = Vec::new();

            for entity in &skinned_entities {
                let render_mesh = match world.get_render_mesh(*entity) {
                    Some(m) => m,
                    None => continue,
                };

                let mesh_id = match self.skinned_meshes.get(&render_mesh.name) {
                    Some(&id) => id,
                    None => continue,
                };

                let skin_index = self
                    .skinning_cache
                    .entity_skin_indices
                    .get(entity)
                    .copied()
                    .unwrap_or(0);
                let joint_offset = self.skinning_cache.get_joint_offset(skin_index);

                skinned_objects.push(SkinnedObjectData {
                    transform_index: 0,
                    mesh_id,
                    material_id: 0,
                    joint_offset,
                });

                let mesh_data = &self.skinned_mesh_data[mesh_id as usize];
                self.skinned_draw_calls.push((
                    mesh_data.index_count,
                    mesh_data.index_offset,
                    skinned_objects.len() as u32 - 1,
                ));
            }

            if !skinned_objects.is_empty() {
                if skinned_objects.len() > self.skinned_object_buffer_size {
                    let new_size = (skinned_objects.len() as f32 * 2.0).ceil() as usize;
                    self.skinned_object_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("Selection Mask Skinned Object Buffer (Resized)"),
                        size: (std::mem::size_of::<SkinnedObjectData>() * new_size) as u64,
                        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                        mapped_at_creation: false,
                    });
                    self.skinned_object_buffer_size = new_size;
                    self.rebuild_skinned_instance_bind_group(device);
                }

                queue.write_buffer(
                    &self.skinned_object_buffer,
                    0,
                    bytemuck::cast_slice(&skinned_objects),
                );
            }
        }

        if self.draw_calls.is_empty() && self.skinned_draw_calls.is_empty() {
            self.enabled = false;
        }
    }

    fn execute<'r, 'e>(
        &mut self,
        context: PassExecutionContext<'r, 'e, crate::ecs::world::World>,
    ) -> crate::render::wgpu::rendergraph::Result<
        Vec<crate::render::wgpu::rendergraph::SubGraphRunCommand<'r>>,
    > {
        let (mask_view, mask_load_op, mask_store_op) =
            context.get_color_attachment("selection_mask")?;

        if !self.enabled || (self.draw_calls.is_empty() && self.skinned_draw_calls.is_empty()) {
            let render_pass = context
                .encoder
                .begin_render_pass(&wgpu::RenderPassDescriptor {
                    label: Some("Selection Mask Clear Pass"),
                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                        view: mask_view,
                        resolve_target: None,
                        ops: wgpu::Operations {
                            load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                            store: mask_store_op,
                        },
                        depth_slice: None,
                    })],
                    depth_stencil_attachment: None,
                    timestamp_writes: None,
                    occlusion_query_set: None,
                });
            drop(render_pass);
            return Ok(context.into_sub_graph_commands());
        }

        if self.total_joints_to_dispatch > 0 && !self.skinned_draw_calls.is_empty() {
            let mut compute_pass =
                context
                    .encoder
                    .begin_compute_pass(&wgpu::ComputePassDescriptor {
                        label: Some("Selection Mask Skinning Compute Pass"),
                        timestamp_writes: None,
                    });

            compute_pass.set_pipeline(&self.skinning_compute_pipeline);
            compute_pass.set_bind_group(0, &self.skinning_compute_bind_group, &[]);

            let num_workgroups = self.total_joints_to_dispatch.div_ceil(64);
            compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
        }

        let (depth_view, depth_load_op, depth_store_op) = context.get_depth_attachment("depth")?;

        let mut render_pass = context
            .encoder
            .begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Selection Mask Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: mask_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: mask_load_op,
                        store: mask_store_op,
                    },
                    depth_slice: None,
                })],
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                    view: depth_view,
                    depth_ops: Some(wgpu::Operations {
                        load: depth_load_op,
                        store: depth_store_op,
                    }),
                    stencil_ops: None,
                }),
                timestamp_writes: None,
                occlusion_query_set: None,
            });

        if !self.draw_calls.is_empty() {
            render_pass.set_pipeline(&self.pipeline);
            render_pass.set_bind_group(0, &self.uniform_bind_group, &[]);
            render_pass.set_bind_group(1, &self.instance_bind_group, &[]);
            render_pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
            render_pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint32);

            for (index_count, index_offset, instance_index) in &self.draw_calls {
                render_pass.draw_indexed(
                    *index_offset..(*index_offset + *index_count),
                    0,
                    *instance_index..(*instance_index + 1),
                );
            }
        }

        if !self.skinned_draw_calls.is_empty() {
            render_pass.set_pipeline(&self.skinned_pipeline);
            render_pass.set_bind_group(0, &self.uniform_bind_group, &[]);
            render_pass.set_bind_group(1, &self.skinned_instance_bind_group, &[]);
            render_pass.set_vertex_buffer(0, self.skinned_vertex_buffer.slice(..));
            render_pass.set_index_buffer(
                self.skinned_index_buffer.slice(..),
                wgpu::IndexFormat::Uint32,
            );

            for (index_count, index_offset, instance_index) in &self.skinned_draw_calls {
                render_pass.draw_indexed(
                    *index_offset..(*index_offset + *index_count),
                    0,
                    *instance_index..(*instance_index + 1),
                );
            }
        }

        drop(render_pass);

        Ok(context.into_sub_graph_commands())
    }
}
