use super::{MeshHandle, MeshPool, TextureHandle, TexturePool};
use std::collections::HashMap;

#[derive(Copy, Clone, PartialEq)]
pub struct MeshInstance {
    pub mesh: MeshHandle,
    pub albedo: TextureHandle,
    pub transform: glam::Mat4,
}

pub struct MeshCommands {
    transforms: Vec<[f32; 16]>,
    draws: Vec<(usize, usize, std::ops::Range<u32>)>,
}

impl MeshCommands {
    pub fn build(instances: &Vec<MeshInstance>) -> Self {
        let mut sorted_meshes = HashMap::<(usize, usize), Vec<MeshInstance>>::new();
        for instance in instances.iter() {
            let group_id = instance.mesh.group_id;
            let texture_id = instance.albedo.id;
            let key = (group_id, texture_id);
            if let Some(by_group) = sorted_meshes.get_mut(&key) {
                by_group.push(*instance);
            } else {
                sorted_meshes.insert(key, vec![*instance]);
            }
        }

        let mut transforms = Vec::new();
        let mut draws = Vec::new();

        // this code assumes MeshHandle only uses group_id (which it does for now)
        // TODO bucket by sub_id too before MeshHandle supports it
        for ((group_id, texture_id), meshes) in sorted_meshes.iter() {
            let start_idx = transforms.len() as u32;
            let as_arrays = meshes.iter().map(|i| i.transform.to_cols_array());
            transforms.extend(as_arrays);
            let end_idx = transforms.len() as u32;
            draws.push((*group_id, *texture_id, start_idx..end_idx));
        }

        Self { transforms, draws }
    }

    pub fn get_storage(&self) -> &[u8] {
        bytemuck::cast_slice(&self.transforms)
    }

    pub fn dispatch<'a>(
        &self,
        rp: &mut wgpu::RenderPass<'a>,
        mesh_pool: &'a MeshPool,
        texture_pool: &'a TexturePool,
    ) {
        // TODO one group per mesh, still...
        // TODO this could be implemented without accessing private members
        for (group_id, texture_id, meshes_range) in self.draws.iter() {
            let group = mesh_pool.groups.get(*group_id).unwrap();
            let texture = texture_pool.textures.get(*texture_id).unwrap();
            rp.set_vertex_buffer(0, group.vertices.slice(..));
            rp.set_index_buffer(group.indices.slice(..), wgpu::IndexFormat::Uint32);
            rp.set_bind_group(2, &texture.bind_group, &[]);
            let indices = 0..(group.index_capacity as u32);
            rp.draw_indexed(indices, 0, meshes_range.to_owned());
        }
    }
}
