pittore 0.2.4

use std::{
    any::TypeId,
    collections::hash_map::Entry,
    ops::{Deref, DerefMut, Range},
};

use rustc_hash::FxHashMap;
use wgpu;

pub trait BufferItem: bytemuck::NoUninit + bytemuck::Pod + wgpu::WasmNotSend {}
impl<T: bytemuck::NoUninit + bytemuck::Pod + wgpu::WasmNotSend> BufferItem for T {}

pub type TypeIdMap<T> = FxHashMap<TypeId, T>;

#[derive(Default)]
pub struct Buffers {
    buffers: TypeIdMap<DynamicBuffer>,
}

impl Buffers {
    pub fn contains_key<T: 'static>(&self) -> bool {
        self.buffers.contains_key(&TypeId::of::<T>())
    }

    pub fn get<T: 'static>(&self) -> Option<&DynamicBuffer> {
        self.buffers.get(&TypeId::of::<T>())
    }

    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut DynamicBuffer> {
        self.buffers.get_mut(&TypeId::of::<T>())
    }

    pub fn entry<T: 'static>(&mut self) -> Entry<TypeId, DynamicBuffer> {
        self.buffers.entry(TypeId::of::<T>())
    }

    pub fn insert(&mut self, buffer: DynamicBuffer) {
        self.buffers.insert(buffer.type_id, buffer);
    }

    pub fn write_all(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        for buffer in self.buffers.values_mut() {
            buffer.write(device, queue);
        }
    }

    pub fn clear_temporary_all(&mut self) {
        for buffer in self.buffers.values_mut() {
            buffer.clear_temporary();
        }
    }
}

pub struct DynamicBuffer {
    pub buffer: wgpu::Buffer,
    pub usage: wgpu::BufferUsages,
    data: Vec<u8>,
    data_temporary: Vec<u8>,
    len_persist_written: usize,
    type_id: TypeId,
    type_size: usize,
    label: String,
}

impl DynamicBuffer {
    /// Create a new dynamic buffer.
    pub fn new<T: BufferItem>(device: &wgpu::Device, usage: wgpu::BufferUsages) -> Self {
        Self::with_capacity::<T>(device, usage, 1)
    }

    /// Create a new dynamic buffer with reserved capacity.
    pub fn with_capacity<T: BufferItem>(
        device: &wgpu::Device,
        usage: wgpu::BufferUsages,
        capacity: usize,
    ) -> Self {
        if size_of::<T>() % 4 != 0 {
            panic!("size of item must be a multiple of 4");
        }

        let label = format!("dynamic buffer for {}", std::any::type_name::<T>());

        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(&label),
            size: (size_of::<T>() * capacity) as u64,
            usage,
            mapped_at_creation: false,
        });

        DynamicBuffer {
            buffer,
            usage,
            data: Vec::new(),
            data_temporary: Vec::new(),
            len_persist_written: 0,
            type_id: TypeId::of::<T>(),
            type_size: size_of::<T>(),
            label,
        }
    }

    /// Append values to the buffer.
    ///
    /// Note that this operation is not immediately executed. Call [`DynamicBuffer::write`] to apply changes.
    pub fn append<T: BufferItem>(&mut self, values: Vec<T>) -> Range<u32> {
        debug_assert_eq!(TypeId::of::<T>(), self.type_id);
        let start = self.data.len() / self.type_size;
        let end = start + values.len();
        let mut values = bytemuck::cast_slice(&values).to_vec();
        self.data.append(&mut values);

        start as u32..end as u32
    }

    /// Append temporary values to the buffer. Temporary values are cleared every frame after rendering.
    ///
    /// Note that this operation is not immediately executed. Call [`DynamicBuffer::write`] to apply changes.
    pub fn append_temporary<T: BufferItem>(&mut self, values: Vec<T>) -> Range<u32> {
        debug_assert_eq!(TypeId::of::<T>(), self.type_id);
        let start = self.data_temporary.len() / self.type_size;
        let end = start + values.len();
        let mut values = bytemuck::cast_slice(&values).to_vec();
        self.data_temporary.append(&mut values);

        start as u32..end as u32
    }

    /// Clear all values and set the length to zero.
    pub fn clear(&mut self) {
        self.data.clear();
        self.len_persist_written = 0;
    }

    /// Clear all temporary values.
    pub fn clear_temporary(&mut self) {
        self.data_temporary.clear();
    }

    /// Write staging changes to the buffer.
    ///
    /// Returns true if new buffer was allocated.
    pub fn write(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) -> bool {
        let len_persist = self.data.len() / self.type_size;
        let len_temporary = self.data_temporary.len() / self.type_size;
        let len_total = len_persist + len_temporary;
        let mut cap = self.buffer.size() as usize / self.type_size;
        self.data.append(&mut self.data_temporary);

        let mut realloced = false;

        if len_total > self.len_persist_written {
            if len_total > cap {
                cap = (cap * 2).max(len_total);

                self.buffer = device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some(&self.label),
                    size: (self.type_size * cap) as u64,
                    usage: self.usage,
                    mapped_at_creation: true,
                });

                self.buffer
                    .slice(..self.data.len() as u64)
                    .get_mapped_range_mut()
                    .copy_from_slice(&self.data);
                self.buffer.unmap();

                realloced = true;
            } else {
                queue.write_buffer(
                    &self.buffer,
                    (self.type_size * self.len_persist_written) as u64,
                    bytemuck::cast_slice(
                        &self.data
                            [self.type_size * self.len_persist_written..self.type_size * len_total],
                    ),
                );
            }
        }

        self.len_persist_written = len_persist;
        self.data.truncate(len_persist);

        realloced
    }

    /// Get the offset for temporary values.
    pub fn offset_temporary(&self) -> u32 {
        self.len_persist_written as u32
    }

    /// Offset the range of temporary values.
    pub fn offset_temporary_range(&self, bounds: Range<u32>) -> Range<u32> {
        let len = self.len_persist_written as u32;
        bounds.start + len..bounds.end + len
    }
}

impl Deref for DynamicBuffer {
    type Target = wgpu::Buffer;

    fn deref(&self) -> &Self::Target {
        &self.buffer
    }
}

impl DerefMut for DynamicBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.buffer
    }
}