//! Chunked GPU buffer management for datasets exceeding `max_storage_buffer_binding_size`.

use super::{GpuContext, GpuError};

/// A buffer that automatically chunks data to respect GPU buffer size limits.
///
/// wgpu's `max_storage_buffer_binding_size` defaults to 128 MiB. At 128d×f32,
/// that's ~67M embeddings per chunk — well above typical use. ChunkedBuffer
/// handles the edge case transparently.
pub struct ChunkedBuffer {
    pub buffers: Vec<wgpu::Buffer>,
    pub chunk_sizes: Vec<usize>, // number of items per chunk
    pub total_items: usize,
    pub item_stride: usize, // bytes per item
}

impl ChunkedBuffer {
    /// Create a chunked storage buffer from flat f32 data.
    ///
    /// `data`: flat f32 slice (n_items × dim)
    /// `dim`: embedding dimension
    /// `label`: debug label for the buffer
    pub fn from_f32(
        ctx: &GpuContext,
        data: &[f32],
        dim: usize,
        label: &str,
        usage: wgpu::BufferUsages,
    ) -> Result<Self, GpuError> {
        let item_stride = dim * std::mem::size_of::<f32>();
        let n_items = data.len() / dim;
        let max_bytes = ctx.limits.max_storage_buffer_binding_size as usize;
        let items_per_chunk = max_bytes / item_stride;

        let mut buffers = Vec::new();
        let mut chunk_sizes = Vec::new();
        let mut offset = 0;

        while offset < n_items {
            let chunk_items = (n_items - offset).min(items_per_chunk);
            let start = offset * dim;
            let end = start + chunk_items * dim;
            let chunk_data = &data[start..end];

            let buffer = ctx.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(&format!("{label}_chunk_{}", buffers.len())),
                contents: bytemuck::cast_slice(chunk_data),
                usage,
            });

            buffers.push(buffer);
            chunk_sizes.push(chunk_items);
            offset += chunk_items;
        }

        Ok(ChunkedBuffer {
            buffers,
            chunk_sizes,
            total_items: n_items,
            item_stride,
        })
    }

    /// Create a single storage buffer from a u32 slice.
    pub fn from_u32(
        ctx: &GpuContext,
        data: &[u32],
        label: &str,
        usage: wgpu::BufferUsages,
    ) -> Result<wgpu::Buffer, GpuError> {
        let byte_size = (data.len() * std::mem::size_of::<u32>()) as u64;
        let max = ctx.limits.max_storage_buffer_binding_size as u64;
        if byte_size > max {
            return Err(GpuError::BufferTooLarge {
                requested: byte_size,
                max,
            });
        }
        Ok(ctx.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some(label),
            contents: bytemuck::cast_slice(data),
            usage,
        }))
    }

    /// Number of chunks.
    pub fn num_chunks(&self) -> usize {
        self.buffers.len()
    }
}

/// Convenience trait for buffer initialization (re-exported from wgpu::util).
pub(crate) use wgpu::util::DeviceExt;
