pub mod mesh;

use bytemuck::Pod;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;

pub struct GpuVec<T> {
    device: Arc<wgpu::Device>,
    data: Vec<T>,
    buffer: wgpu::Buffer,
    capacity: usize,
    usage: wgpu::BufferUsages,
    label: Option<String>,
    realign: Option<usize>,
}

impl<T: Clone + Pod> GpuVec<T> {
    pub fn new(
        device: Arc<wgpu::Device>,
        usage: wgpu::BufferUsages,
        initial_capacity: usize,
        label: Option<String>,
        uses_dynamic_offsets: bool,
    ) -> Self {
        let capacity_bytes = initial_capacity * std::mem::size_of::<T>();
        let usage = usage | wgpu::BufferUsages::COPY_DST;
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: label.as_ref().map(|s| s.as_str()),
            size: capacity_bytes as wgpu::BufferAddress,
            usage,
            mapped_at_creation: false,
        });

        let data = Vec::with_capacity(initial_capacity);

        let realign = if uses_dynamic_offsets {
            if usage.contains(wgpu::BufferUsages::STORAGE) {
                Some(device.limits().min_storage_buffer_offset_alignment as usize)
            } else if usage.contains(wgpu::BufferUsages::UNIFORM) {
                Some(device.limits().min_uniform_buffer_offset_alignment as usize)
            } else {
                None
            }
        } else {
            None
        };

        Self {
            device,
            data,
            buffer,
            capacity: capacity_bytes,
            usage,
            label,
            realign,
        }
    }

    pub fn write(&mut self, queue: &wgpu::Queue) -> &wgpu::Buffer {
        let capacity_bytes = self.buf_offset(self.data.capacity());
        if capacity_bytes != self.capacity && capacity_bytes != 0 {
            self.capacity = capacity_bytes;
            self.buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
                label: self.label.as_ref().map(|s| s.as_str()),
                size: capacity_bytes as wgpu::BufferAddress,
                usage: self.usage,
                mapped_at_creation: false,
            });
        }

        let mut realigned = Vec::new();
        let src_buf = if let Some(realign) = self.realign {
            realigned.resize(self.data.len() * realign, 0);

            for (index, elem) in self.data.iter().enumerate() {
                let elem = [elem.clone()];
                let casted: &[u8] = bytemuck::cast_slice(&elem);
                let dst_offset = index * realign;
                realigned[dst_offset..(dst_offset + casted.len())].copy_from_slice(casted);
            }

            realigned.as_slice()
        } else {
            bytemuck::cast_slice(self.data.as_slice())
        };

        queue.write_buffer(&self.buffer, 0, src_buf);

        &self.buffer
    }

    pub fn buf_offset(&self, index: usize) -> usize {
        if let Some(realign) = self.realign {
            index * realign
        } else {
            index * std::mem::size_of::<T>()
        }
    }
}

impl<T: Pod> AsRef<wgpu::Buffer> for GpuVec<T> {
    fn as_ref(&self) -> &wgpu::Buffer {
        &self.buffer
    }
}

impl<T: Pod> Deref for GpuVec<T> {
    type Target = Vec<T>;

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

impl<T: Pod> DerefMut for GpuVec<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.data
    }
}