twgpu 0.4.1

use std::collections::HashMap;
use twmap::{Envelope, Layer, Quad, TilesLayer, TwMap};
use vek::az::UnwrappedAs;
use vek::Rgba;
use wgpu::{
    BindGroupLayoutEntry, BindingType, Device, Extent3d, Queue, ShaderStages,
    TexelCopyBufferLayout, Texture, TextureDescriptor, TextureDimension, TextureFormat,
    TextureSampleType, TextureUsages, TextureView, TextureViewDescriptor, TextureViewDimension,
};

use super::EnvelopeSampling;

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

pub struct GpuEnvelopesData {
    pub index: GpuEnvelopesIndex,
    pub texture: Texture,
}

/// To calculate and update only the envelope values we need, we store each index + offset combination that is used in the map
/// We put those unique combinations into a list to calculate the envelope values from it
/// In the list calculated for the gpu buffer, the first two elements are the neutral elements for the color/position.
/// Note that there must be less than i32::MAX - 2 envelope points.
#[derive(Debug, Clone, Default)]
pub struct GpuEnvelopesIndex {
    // TODO: Change i32 to u32 as soon as naga supports textureLoad with u32
    /// Mapping of index + offset to internal envelope value list index
    pub mapping: HashMap<(u16, i32), i32>,
    /// Internal list of envelope 'values'
    /// The values are stored only in gpu memory only, calculated using this list
    pub ordered: Vec<(u16, i32)>,
}

pub const NEUTRAL_COLOR_INDEX: i32 = 0;
pub const NEUTRAL_POSITION_INDEX: i32 = 1;

impl GpuEnvelopesIndex {
    /// Add a potentially new index + offset combination
    pub fn register(&mut self, index: Option<u16>, offset: i32) {
        let new_potential_env_index = self.size().unwrapped_as();
        if let Some(identifier) = index.map(|n| (n, offset)) {
            self.mapping.entry(identifier).or_insert_with(|| {
                self.ordered.push(identifier);
                new_potential_env_index
            });
        }
    }

    /// Registers all index + offset combinations from the map, excluding sound envelopes
    pub fn new(map: &TwMap) -> Self {
        let mut index = GpuEnvelopesIndex::default();
        for layer in map.groups.iter().flat_map(|g| &g.layers) {
            match layer {
                Layer::Tiles(l) => index.register(l.color_env, l.color_env_offset),
                Layer::Quads(l) => {
                    for quad in &l.quads {
                        index.register(quad.color_env, quad.color_env_offset);
                        index.register(quad.position_env, quad.position_env_offset);
                    }
                }
                _ => {}
            }
        }
        index
    }

    /// Returns the amount of elements in the index, plus the implicit neutral elements
    pub fn size(&self) -> usize {
        self.ordered.len() + 2
    }
}

impl GpuEnvelopesData {
    pub fn layout_entry(binding: u32) -> BindGroupLayoutEntry {
        BindGroupLayoutEntry {
            binding,
            visibility: ShaderStages::VERTEX,
            ty: BindingType::Texture {
                sample_type: TextureSampleType::Float { filterable: false },
                view_dimension: TextureViewDimension::D1,
                multisampled: false,
            },
            count: None,
        }
    }

    pub fn view(&self) -> TextureView {
        self.texture.create_view(&TextureViewDescriptor {
            dimension: Some(TextureViewDimension::D1),
            ..TextureViewDescriptor::default()
        })
    }

    pub fn quad_color_env_index(&self, quad: &Quad) -> i32 {
        match quad.color_env {
            Some(n) => *self.index.mapping.get(&(n, quad.color_env_offset)).unwrap(),
            None => NEUTRAL_COLOR_INDEX,
        }
    }

    pub fn quad_position_env_index(&self, quad: &Quad) -> i32 {
        match quad.position_env {
            Some(n) => *self
                .index
                .mapping
                .get(&(n, quad.position_env_offset))
                .unwrap(),
            None => NEUTRAL_POSITION_INDEX,
        }
    }

    pub fn tiles_color_env_index(&self, layer: &TilesLayer) -> i32 {
        match layer.color_env {
            Some(n) => *self
                .index
                .mapping
                .get(&(n, layer.color_env_offset))
                .unwrap(),
            None => NEUTRAL_COLOR_INDEX,
        }
    }

    pub fn upload(map: &TwMap, device: &Device) -> Self {
        let index = GpuEnvelopesIndex::new(map);
        let mut initial_data = Vec::with_capacity(index.size());
        initial_data.extend([Rgba::white(), Rgba::zero()]);
        for &(i, offset) in &index.ordered {
            let value = map.envelopes[i.unwrapped_as::<usize>()].sample(0, 0, offset);
            initial_data.push(value);
        }
        let texture = device.create_texture(&TextureDescriptor {
            label: LABEL,
            size: Extent3d {
                width: index.size().unwrapped_as(),
                height: 1,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D1,
            format: TextureFormat::Rgba32Float,
            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
            view_formats: &[],
        });
        Self { index, texture }
    }

    pub fn update(
        &self,
        envelopes: &[Envelope],
        client_micros: i64,
        server_micros: i64,
        queue: &Queue,
    ) {
        let mut data = Vec::with_capacity(self.index.size());
        data.extend([Rgba::white(), Rgba::zero()]);
        for &(i, offset) in &self.index.ordered {
            let env = &envelopes[i.unwrapped_as::<usize>()];
            let value = env.sample(client_micros, server_micros, offset);
            data.push(value);
        }
        queue.write_texture(
            self.texture.as_image_copy(),
            bytemuck::cast_slice(&data),
            TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: None,
                rows_per_image: None,
            },
            Extent3d {
                width: self.index.size().unwrapped_as(),
                height: 1,
                depth_or_array_layers: 1,
            },
        );
    }
}