use super::handle::{MeshHandle, TextureHandle};
use super::mesh::MeshData;
use wgpu::util::DeviceExt;

pub struct MeshGroup {
    // TODO make these all private
    pub vertices: wgpu::Buffer,
    pub vertex_capacity: usize,
    pub indices: wgpu::Buffer,
    pub index_capacity: usize,
}

impl MeshGroup {
    fn new(device: &wgpu::Device, data: &MeshData) -> Self {
        let vertex_capacity = data.vertices.len();
        let vertices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Vertex Buffer"),
            contents: bytemuck::cast_slice(&data.vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });

        let index_capacity = data.indices.len();
        let indices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Index Buffer"),
            contents: bytemuck::cast_slice(&data.indices),
            usage: wgpu::BufferUsages::INDEX,
        });

        Self {
            vertex_capacity,
            vertices,
            index_capacity,
            indices,
        }
    }
}

#[derive(Default)]
pub struct MeshPool {
    // TODO make this private
    pub groups: slab::Slab<MeshGroup>,
}

impl MeshPool {
    pub fn allocate(&mut self, device: &wgpu::Device, data: &MeshData) -> MeshHandle {
        let group = MeshGroup::new(device, data);
        let group_id = self.groups.insert(group);
        let sub_id = 0;
        MeshHandle { group_id, sub_id }
    }

    pub fn get_group(&self, handle: &MeshHandle) -> Option<&MeshGroup> {
        self.groups.get(handle.group_id)
    }
}

pub struct TextureData {
    pub width: u32,
    pub height: u32,
    pub data: Vec<u8>,
}

pub struct Texture {
    // TODO make this all private
    pub texture: wgpu::Texture,
    pub bind_group: wgpu::BindGroup,
}

impl Texture {
    pub fn new(
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        sampler: &wgpu::Sampler,
        bind_group_layout: &wgpu::BindGroupLayout,
        data: &TextureData,
    ) -> Self {
        let size = wgpu::Extent3d {
            width: data.width,
            height: data.height,
            depth_or_array_layers: 1,
        };

        let texture = device.create_texture(&wgpu::TextureDescriptor {
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            label: None,
        });

        let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());

        queue.write_texture(
            wgpu::ImageCopyTexture {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &data.data,
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: std::num::NonZeroU32::new(4 * size.width),
                rows_per_image: std::num::NonZeroU32::new(size.height),
            },
            size,
        );

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&texture_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
            label: None,
        });

        Texture {
            texture,
            bind_group,
        }
    }
}

pub struct TexturePool {
    // TODO make this all private
    pub textures: slab::Slab<Texture>,
    pub sampler: wgpu::Sampler,
    pub bind_group_layout: wgpu::BindGroupLayout,
}

impl TexturePool {
    pub fn new(device: &wgpu::Device) -> Self {
        let textures = Default::default();
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        multisampled: false,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
            label: Some("Texture Bind Group Layout"),
        });

        Self {
            textures,
            sampler,
            bind_group_layout,
        }
    }

    pub fn allocate(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        data: &TextureData,
    ) -> TextureHandle {
        let texture = Texture::new(device, queue, &self.sampler, &self.bind_group_layout, data);
        let id = self.textures.insert(texture);
        TextureHandle { id }
    }
}
