cgpu 0.1.0

use std::collections::HashMap;

use crate::error::{GpuError, GpuFailure};

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum BufferUse {
    Input,
    Output,
    Scratch,
    Staging,
    Metadata,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BufferLease {
    pub id: u64,
    pub bytes: usize,
    pub use_case: BufferUse,
}

#[derive(Clone, Debug)]
struct BufferSlot {
    lease: BufferLease,
    in_use: bool,
}

#[derive(Clone, Debug, Default)]
pub struct BufferPool {
    next_id: u64,
    slots: Vec<BufferSlot>,
}

impl BufferPool {
    pub fn acquire(&mut self, bytes: usize, use_case: BufferUse) -> BufferLease {
        for slot in &mut self.slots {
            if !slot.in_use && slot.lease.bytes >= bytes && slot.lease.use_case == use_case {
                slot.in_use = true;
                return slot.lease;
            }
        }

        let lease = BufferLease {
            id: self.next_id,
            bytes,
            use_case,
        };
        self.next_id = self.next_id.saturating_add(1);
        self.slots.push(BufferSlot {
            lease,
            in_use: true,
        });
        lease
    }

    pub fn release(&mut self, lease: BufferLease) -> bool {
        if let Some(slot) = self.slots.iter_mut().find(|slot| slot.lease.id == lease.id) {
            slot.in_use = false;
            return true;
        }

        false
    }

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

#[derive(Clone, Debug)]
pub struct ResourceCache<T = ResourceEntry> {
    entries: HashMap<String, T>,
}

impl<T> Default for ResourceCache<T> {
    fn default() -> Self {
        Self {
            entries: HashMap::new(),
        }
    }
}

impl<T> ResourceCache<T> {
    pub fn get_or_insert_with(&mut self, key: impl Into<String>, make: impl FnOnce() -> T) -> &T {
        let key = key.into();
        self.entries.entry(key).or_insert_with(make)
    }

    pub fn insert(&mut self, key: impl Into<String>, entry: T) {
        self.entries.insert(key.into(), entry);
    }

    pub fn get(&self, key: &str) -> Option<&T> {
        self.entries.get(key)
    }

    pub fn contains_key(&self, key: &str) -> bool {
        self.entries.contains_key(key)
    }

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

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

    pub fn clear(&mut self) {
        self.entries.clear();
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ResourceEntry {
    pub label: String,
    pub bytes: usize,
}

#[cfg(feature = "gpu")]
pub fn storage_binding(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

#[cfg(feature = "gpu")]
pub fn create_mapped_buffer(
    device: &wgpu::Device,
    label: &str,
    bytes: &[u8],
    usage: wgpu::BufferUsages,
) -> wgpu::Buffer {
    let buffer_size =
        u64::try_from(bytes.len().max(std::mem::size_of::<u32>())).unwrap_or(u64::MAX);
    let buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size: buffer_size,
        usage,
        mapped_at_creation: true,
    });
    if !bytes.is_empty() {
        let mut mapped = buffer.slice(..).get_mapped_range_mut();
        mapped.slice(0..bytes.len()).copy_from_slice(bytes);
    }
    buffer.unmap();
    buffer
}

#[cfg(feature = "gpu")]
pub fn create_zeroed_buffer(
    device: &wgpu::Device,
    label: &str,
    size: u64,
    usage: wgpu::BufferUsages,
) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size: size.max(u64::try_from(std::mem::size_of::<u32>()).unwrap()),
        usage,
        mapped_at_creation: false,
    })
}

pub fn u32_slice_bytes(values: &[u32]) -> Vec<u8> {
    let mut bytes = Vec::with_capacity(std::mem::size_of_val(values));
    for value in values {
        bytes.extend_from_slice(&value.to_ne_bytes());
    }
    bytes
}

pub fn read_u32(bytes: &[u8], offset: usize) -> Result<u32, GpuError> {
    let value = bytes
        .get(offset..offset + std::mem::size_of::<u32>())
        .ok_or(GpuError::GpuFailure(GpuFailure::Readback))?;
    Ok(u32::from_ne_bytes(
        value
            .try_into()
            .map_err(|_| GpuError::GpuFailure(GpuFailure::Readback))?,
    ))
}