use wgpu::{Buffer, Device, Queue};
use wgpu::util::DeviceExt;
use std::mem::size_of;
use crate::geometry_manager::Shape;
use crate::renderer::{InstanceRaw, RenderInstance};

pub struct InstanceManager {
    instances: Vec<RenderInstance>,
    raw: Vec<InstanceRaw>,
    buffer: Buffer,
}

impl InstanceManager {
    pub fn new(device: &Device) -> Self {
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Instance Buffer (empty)"),
            size: 1,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            instances: Vec::new(),
            raw: Vec::new(),
            buffer,
        }
    }

    pub fn set_instances(&mut self, device: &Device, queue: &Queue, instances: Vec<RenderInstance>) {
        self.raw = instances.iter().map(RenderInstance::to_raw).collect();
        let byte_len = (self.raw.len() * size_of::<InstanceRaw>()) as wgpu::BufferAddress;

        if byte_len > self.buffer.size() {
            self.buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("Instance Buffer (resized)"),
                contents: bytemuck::cast_slice(&self.raw),
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            });
        } else {
            queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&self.raw));
        }

        self.instances = instances;
    }

    pub fn raw_instances_for_shape(&self, shape: Shape) -> Vec<&InstanceRaw> {
        self.instances
            .iter()
            .zip(self.raw.iter())
            .filter(|(inst, _)| inst.shape == shape)
            .map(|(_, raw)| raw)
            .collect()
    }

    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    pub fn len(&self) -> usize {
        self.instances.len()
    }

    pub fn is_empty(&self) -> bool {
        self.instances.is_empty()
    }
}