twgpu 0.4.1

use std::ops::Range;

use crate::buffer::{GpuBuffer, GpuDeviceExt, GpuVecBuffer};
use wgpu::util::RenderEncoder;
use wgpu::{
    include_wgsl, BindGroupLayout, BindGroupLayoutDescriptor, BlendState, BufferUsages,
    ColorTargetState, ColorWrites, Device, FragmentState, IndexFormat, PipelineCompilationOptions,
    PipelineLayoutDescriptor, PrimitiveState, PrimitiveTopology, RenderBundle,
    RenderBundleDescriptor, RenderBundleEncoderDescriptor, RenderPipeline,
    RenderPipelineDescriptor, TextureFormat, TextureViewDimension, VertexState,
};

use crate::sprites::texture_storage::SpriteTextures;
use crate::sprites::{SpriteVertex, TextureToken};
use crate::textures::{Mapres, Samplers};
use crate::{Camera, GpuCamera};

const LABEL: Option<&str> = Some("Sprites Static");

pub struct SpritesStatic {
    pub format: TextureFormat,
    pub bind_group_layout: BindGroupLayout,
    pub pipeline: RenderPipeline,
    pub index_buffer: GpuBuffer<[u32; 6]>,
}

impl SpritesStatic {
    pub fn render<'pass, R: RenderEncoder<'pass>>(
        &'pass self,
        camera_index: usize,
        sprite_buffer: &'pass GpuVecBuffer<[SpriteVertex; 4]>,
        sprite_textures: &[TextureToken],
        textures: &'pass SpriteTextures,
        render_pass: &mut R,
    ) {
        self.render_range(
            0..sprite_buffer.len(),
            camera_index,
            sprite_buffer,
            sprite_textures,
            textures,
            render_pass,
        )
    }

    pub fn render_range<'pass, R: RenderEncoder<'pass>>(
        &'pass self,
        range: Range<usize>,
        camera_index: usize,
        sprite_buffer: &'pass GpuVecBuffer<[SpriteVertex; 4]>,
        sprite_textures: &[TextureToken],
        textures: &'pass SpriteTextures,
        render_pass: &mut R,
    ) {
        assert_eq!(sprite_buffer.len(), sprite_textures.len());
        if range.is_empty() {
            return;
        }
        render_pass.set_pipeline(&self.pipeline);
        render_pass.set_index_buffer(self.index_buffer.slice(), IndexFormat::Uint32);

        let mut texture_iter = sprite_textures[range.clone()].iter();
        render_pass.set_vertex_buffer(0, sprite_buffer.full_slice());
        for i in range {
            let texture = texture_iter.next().unwrap();
            let camera_offset = camera_index * std::mem::size_of::<Camera>();
            render_pass.set_bind_group(
                0,
                Some(textures.get_bind_group(*texture)),
                &[camera_offset as u32],
            );
            render_pass.draw_indexed(0..6, i as i32 * 4, 0..1);
        }
    }

    pub fn prepare_render_with_temporary_buffer(
        &self,
        sprites: &[[SpriteVertex; 4]],
        sprite_textures: &[TextureToken],
        textures: &SpriteTextures,
        device: &Device,
    ) -> RenderBundle {
        let mut render_pass = device.create_render_bundle_encoder(&RenderBundleEncoderDescriptor {
            label: LABEL,
            color_formats: &[Some(self.format)],
            depth_stencil: None,
            sample_count: 1,
            multiview: None,
        });
        let vertex_buffer = device.array_buffer(sprites, LABEL, BufferUsages::VERTEX);
        render_pass.set_pipeline(&self.pipeline);
        render_pass.set_index_buffer(self.index_buffer.slice(), IndexFormat::Uint32);
        render_pass.set_vertex_buffer(0, vertex_buffer.full_slice());
        for (i, texture) in sprite_textures.iter().enumerate() {
            render_pass.set_bind_group(0, textures.get_bind_group(*texture), &[]);
            render_pass.draw_indexed(0..6, i as i32 * 4, 0..1);
        }
        render_pass.finish(&RenderBundleDescriptor { label: LABEL })
    }

    pub fn new(format: TextureFormat, device: &Device) -> Self {
        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: LABEL,
            entries: &[
                GpuCamera::bind_group_layout_entry(0, true),
                Mapres::texture_layout_entry(3, TextureViewDimension::D2),
                Samplers::bind_group_layout_entry(4),
            ],
        });
        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: LABEL,
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });
        let shader_module = device.create_shader_module(include_wgsl!("sprites.wgsl"));
        let vertex_state = VertexState {
            module: &shader_module,
            entry_point: Some("vs_main"),
            compilation_options: PipelineCompilationOptions::default(),
            buffers: &[SpriteVertex::vertex_buffer_layout()],
        };
        let fragment_state = FragmentState {
            module: &shader_module,
            entry_point: Some("fs_main"),
            compilation_options: PipelineCompilationOptions::default(),
            targets: &[Some(ColorTargetState {
                format,
                blend: Some(BlendState::ALPHA_BLENDING),
                write_mask: ColorWrites::all(),
            })],
        };
        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: LABEL,
            layout: Some(&pipeline_layout),
            vertex: vertex_state,
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleList,
                ..PrimitiveState::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(fragment_state),
            multiview: None,
            cache: None,
        });
        const INDEX_DATA: [u32; 6] = [0, 1, 3, 0, 2, 3];
        let index_buffer = device.buffer(&INDEX_DATA, LABEL, BufferUsages::INDEX);
        Self {
            format,
            bind_group_layout,
            pipeline,
            index_buffer,
        }
    }
}