bevy_cube_marcher 0.5.1

use std::marker::PhantomData;
use std::num::NonZero;

use crate::*;
use bevy::asset::{RenderAssetUsages, embedded_asset, load_embedded_asset};
use bevy::core_pipeline::schedule::camera_driver;
use bevy::ecs::schedule::ScheduleConfigs;
use bevy::ecs::system::ScheduleSystem;
use bevy::platform::collections::{HashMap, HashSet};
use bevy::render::extract_component::{ExtractComponent, ExtractComponentPlugin};
use bevy::render::extract_resource::{ExtractResource, ExtractResourcePlugin};
use bevy::render::gpu_readback::{Readback, ReadbackComplete};
use bevy::render::render_asset::RenderAssets;
use bevy::render::render_resource::binding_types::{
    storage_buffer, storage_buffer_read_only, storage_buffer_sized, uniform_buffer,
};
use bevy::render::render_resource::{
    BindGroup, BindGroupEntries, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntries,
    BindGroupLayoutEntryBuilder, Buffer, BufferUsages, CachedComputePipelineId,
    ComputePassDescriptor, ComputePipelineDescriptor, PipelineCache, ShaderStages, ShaderType,
    StorageBuffer, UniformBuffer,
};
use bevy::render::renderer::{RenderContext, RenderDevice, RenderQueue};
use bevy::render::storage::{GpuShaderBuffer, ShaderBuffer};
use bevy::render::{MainWorld, Render, RenderApp, RenderSystems};
pub use bevy::shader::ShaderRef;
use encase::internal::WriteInto;

/// Resources shared across all instances.
struct GpuMarchingCubesGlobalPlugin;

impl Plugin for GpuMarchingCubesGlobalPlugin {
    fn build(&self, app: &mut App) {
        {
            embedded_asset!(app, "marching_cubes.wgsl");
        }

        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };
        render_app
            .add_systems(Render, clear_old_dispatches.in_set(RenderSystems::Cleanup))
            .add_systems(RenderGraph, chunk_generator_pass.before(camera_driver))
            .init_resource::<ChunkGeneratorDispatches>();
    }
}

/// The plugin. Add this to your app or the crate won't work!
///
/// - `Sampler` implements [`ChunkComputeShader`].
/// - `ExtraBufferCache` implements [`GpuExtraBufferCache`]. May be `()`.
/// - `Material` is a bevy [`Material`] that will be applied to the chunks.
pub struct GpuMarchingCubesPlugin<Sampler, ExtraBufferCache, Material> {
    _marker: PhantomData<(Sampler, ExtraBufferCache, Material)>,
}

impl<Sampler, ExtraBufferCache, Material> Default
    for GpuMarchingCubesPlugin<Sampler, ExtraBufferCache, Material>
{
    fn default() -> Self {
        Self {
            _marker: PhantomData,
        }
    }
}

impl<
    Sampler: GpuChunkComputer + Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
    Material: Asset + bevy::prelude::Material,
> Plugin for GpuMarchingCubesPlugin<Sampler, ExtraBufferCache, Material>
{
    fn build(&self, app: &mut App) {
        if !app.is_plugin_added::<GpuMarchingCubesGlobalPlugin>() {
            app.add_plugins(GpuMarchingCubesGlobalPlugin);
        }

        app.add_plugins((
            ExtractResourcePlugin::<ChunkGeneratorSettings<Sampler>>::default(),
            ExtractResourcePlugin::<GpuChunkGeneratorSettings<Sampler>>::default(),
            ExtractComponentPlugin::<ChunkRenderData<Sampler>>::default(),
        ))
        .init_resource::<GpuChunkGeneratorSettings<Sampler>>()
        .add_systems(
            Update,
            (
                init_cache::<Sampler>,
                (
                    update_chunk_loaders::<Sampler>,
                    queue_chunks::<Sampler>,
                    start_chunks::<Sampler, ExtraBufferCache, Material>.run_if(
                        |done_loading: If<
                            Res<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>,
                        >| { done_loading.done },
                    ),
                )
                    .chain()
                    .run_if(resource_exists::<ChunkGeneratorCache<Sampler>>),
            )
                .chain()
                .in_set(ChunkGenSystems),
        );

        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };
        render_app
            .init_resource::<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>()
            .init_resource::<GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache>>()
            .add_systems(ExtractSchedule, extract_pipelines_done::<Sampler>)
            .add_systems(
                Render, // RenderStartup
                init_compute_pipelines::<Sampler, ExtraBufferCache>
                    .run_if(resource_exists::<ChunkGeneratorSettings<Sampler>>) // Might not exist on the first frame due to extraction
                    .run_if(not(resource_exists::<
                        ChunkGeneratorComputePipelines<Sampler>,
                    >)),
            )
            .add_systems(
                Render,
                (
                    clear_gpu_cache::<Sampler, ExtraBufferCache>,
                    check_pipelines::<Sampler>,
                    (
                        calc_needed_buffers::<Sampler, ExtraBufferCache>,
                        ExtraBufferCache::create_extra_buffers(),
                        ExtraBufferCache::clear_extra_buffers(),
                        prepare_bind_groups::<Sampler, ExtraBufferCache>,
                    )
                        .chain()
                        .run_if(
                            |pipelines_done_loading: If<
                                Res<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>,
                            >| pipelines_done_loading.done,
                        ),
                )
                    .chain()
                    .in_set(RenderSystems::PrepareBindGroups),
            );
    }
}

const WORKGROUP_SIZE: u32 = 8;

#[derive(Resource, Default)]
struct ChunkGeneratorDispatches {
    dispatches: Vec<GpuBufferCache>,
}

fn chunk_generator_pass(mut render_context: RenderContext, world: &World) -> Result {
    let dispatches = world.resource::<ChunkGeneratorDispatches>();
    let pipeline_cache = world.resource::<PipelineCache>();
    trace!("Rendering {} chunks", dispatches.dispatches.len());

    for dispatch in dispatches.dispatches.iter() {
        let sample_pipeline = pipeline_cache
            .get_compute_pipeline(dispatch.sample_pipeline)
            .expect("sample pipeline wasn't finished generating");
        let march_pipeline = pipeline_cache
            .get_compute_pipeline(dispatch.march_pipeline)
            .expect("march pipeline wasn't finished generating");

        {
            let mut pass =
                render_context
                    .command_encoder()
                    .begin_compute_pass(&ComputePassDescriptor {
                        label: Some("marching cubes sample pass"),
                        ..default()
                    });

            pass.set_bind_group(0, &dispatch.sample_bind_group, &[]);
            pass.set_pipeline(sample_pipeline);
            pass.dispatch_workgroups(
                dispatch.sample_workgroups,
                dispatch.sample_workgroups,
                dispatch.sample_workgroups,
            );
        }

        {
            let mut pass =
                render_context
                    .command_encoder()
                    .begin_compute_pass(&ComputePassDescriptor {
                        label: Some("marching cubes march pass"),
                        ..default()
                    });

            pass.set_bind_group(0, &dispatch.march_bind_group, &[]);
            pass.set_pipeline(march_pipeline);
            pass.dispatch_workgroups(
                dispatch.march_workgroups,
                dispatch.march_workgroups,
                dispatch.march_workgroups,
            );
        }
    }

    Ok(())
}

#[derive(Resource)]
struct ChunkGeneratorComputePipelines<Sampler> {
    sample_layout: BindGroupLayoutDescriptor,
    sample_pipeline: CachedComputePipelineId,
    march_layout: BindGroupLayoutDescriptor,
    march_pipeline: CachedComputePipelineId,
    _marker: std::marker::PhantomData<Sampler>,
}

fn init_compute_pipelines<
    Sampler: GpuChunkComputer + Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
>(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    pipeline_cache: Res<PipelineCache>,
    settings: Res<ChunkGeneratorSettings<Sampler>>,
    mut pipelines_done_loading: ResMut<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>,
) -> Result<()> {
    trace!("Init compute pipelines");

    pipelines_done_loading.done = false;

    let sample_layout = BindGroupLayoutDescriptor::new(
        "marching cubes sample bind group layout",
        &[
            uniform_buffer::<IVec3>(false),
            uniform_buffer::<MeshSettings>(false),
            storage_buffer::<Vec<f32>>(false),
        ]
        .into_iter()
        .chain(ExtraBufferCache::define_extra_buffers())
        .enumerate()
        .map(|(i, b)| b.build(i as u32, ShaderStages::COMPUTE))
        .collect::<Vec<_>>(),
    );
    let sample_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
        label: Some("marching cubes sample compute shader".into()),
        layout: vec![sample_layout.clone()],
        shader: match Sampler::shader() {
            ShaderRef::Default => {
                return Err(format!("{} shader was not given", ShortName::of::<Sampler>()).into());
            }
            ShaderRef::Handle(handle) => handle,
            ShaderRef::Path(path) => asset_server.load(path),
        },
        ..default()
    });

    let march_layout = BindGroupLayoutDescriptor::new(
        "marching cubes march bind group layout",
        &BindGroupLayoutEntries::sequential(
            ShaderStages::COMPUTE,
            (
                storage_buffer_read_only::<Vec<f32>>(false),
                uniform_buffer::<MeshSettings>(false),
                storage_buffer_sized(false, Some(settings.vertices_buffer_size())),
                storage_buffer::<u32>(false),
                storage_buffer_sized(false, Some(settings.triangles_buffer_size())),
                storage_buffer::<u32>(false),
            ),
        ),
    );
    let march_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
        label: Some("marching cubes march compute shader".into()),
        layout: vec![march_layout.clone()],
        shader: load_embedded_asset!(asset_server.as_ref(), "marching_cubes.wgsl"),
        entry_point: Some("main".into()),
        ..default()
    });

    commands.insert_resource(ChunkGeneratorComputePipelines::<Sampler> {
        sample_layout,
        sample_pipeline,
        march_layout,
        march_pipeline,
        _marker: default(),
    });

    Ok(())
}

/// Reference to the shader used for sampling the density field.
pub trait GpuChunkComputer {
    fn shader() -> ShaderRef;
}

/// Holds the buffers for a chunk being generated.
#[derive(Component, ExtractComponent, Debug)]
pub struct ChunkRenderData<Sampler: Send + Sync + 'static> {
    position: IVec3,
    buffers: BufferCache,
    num_readbacks_done: usize,
    _marker: std::marker::PhantomData<Sampler>,
}

impl<Sampler: Send + Sync + 'static> ChunkRenderData<Sampler> {
    pub fn buffers(&self) -> &BufferCache {
        &self.buffers
    }
}

impl<Sampler: Send + Sync + 'static> Clone for ChunkRenderData<Sampler> {
    fn clone(&self) -> Self {
        Self {
            position: self.position,
            buffers: self.buffers.clone(),
            num_readbacks_done: self.num_readbacks_done,
            _marker: self._marker,
        }
    }
}

/// Holds the main vertex/triangles buffers. May be hashed and used as an identifying key for a loading chunk.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct BufferCache {
    vertices: Handle<ShaderBuffer>,
    num_vertices: Handle<ShaderBuffer>,
    triangles: Handle<ShaderBuffer>,
    num_triangles: Handle<ShaderBuffer>,
}

impl BufferCache {
    fn new(
        vertices_buffer_size: NonZero<u64>,
        triangles_buffer_size: NonZero<u64>,
        buffers: &mut Assets<ShaderBuffer>,
    ) -> Self {
        let vertices_buffer_size: u64 = vertices_buffer_size.into();
        let mut vertices =
            ShaderBuffer::with_size(vertices_buffer_size as usize, RenderAssetUsages::default());
        vertices.buffer_description.usage |= BufferUsages::COPY_SRC;

        let mut num_vertices =
            ShaderBuffer::with_size(size_of::<u32>(), RenderAssetUsages::default());
        num_vertices.buffer_description.usage |= BufferUsages::COPY_SRC | BufferUsages::COPY_DST;

        let triangles_buffer_size: u64 = triangles_buffer_size.into();
        let mut triangles =
            ShaderBuffer::with_size(triangles_buffer_size as usize, RenderAssetUsages::default());
        triangles.buffer_description.usage |= BufferUsages::COPY_SRC;

        let mut num_triangles =
            ShaderBuffer::with_size(size_of::<u32>(), RenderAssetUsages::default());
        num_triangles.buffer_description.usage |= BufferUsages::COPY_SRC | BufferUsages::COPY_DST;

        Self {
            vertices: buffers.add(vertices),
            num_vertices: buffers.add(num_vertices),
            triangles: buffers.add(triangles),
            num_triangles: buffers.add(num_triangles),
        }
    }
}

#[derive(Debug, Clone)]
enum BufferCacheAvailable {
    Available,
    Unavailable,
}

/// Defines extra buffers to be used in the density sampler.
pub trait GpuExtraBufferCache {
    fn define_extra_buffers() -> Vec<BindGroupLayoutEntryBuilder>;
    fn create_extra_buffers() -> ScheduleConfigs<ScheduleSystem>;
    fn clear_extra_buffers() -> ScheduleConfigs<ScheduleSystem>;
    fn buffers(&self) -> Vec<Buffer>;
    /// The number of times [ReadbackReallyComplete] will be triggered by extra buffers
    fn num_extra_readbacks() -> usize;
}

impl GpuExtraBufferCache for () {
    fn define_extra_buffers() -> Vec<BindGroupLayoutEntryBuilder> {
        vec![]
    }
    fn create_extra_buffers() -> ScheduleConfigs<ScheduleSystem> {
        IntoSystem::into_system(|| {}).into_configs()
    }
    fn clear_extra_buffers() -> ScheduleConfigs<ScheduleSystem> {
        IntoSystem::into_system(|| {}).into_configs()
    }
    fn buffers(&self) -> Vec<Buffer> {
        vec![]
    }
    fn num_extra_readbacks() -> usize {
        0
    }
}

/// Holds extra buffers to be used when sampling the density field. Cpu side.
#[derive(Resource, Debug, Clone)]
pub struct GpuCpuChunkGeneratorCache<Sampler> {
    buffer_cache_availability: HashMap<BufferCache, BufferCacheAvailable>,
    _marker: std::marker::PhantomData<Sampler>,
}

impl<Sampler: Send + Sync> GpuCpuChunkGeneratorCache<Sampler> {
    fn new(settings: &ChunkGeneratorSettings<Sampler>, buffers: &mut Assets<ShaderBuffer>) -> Self {
        Self {
            buffer_cache_availability: (0..settings.num_buffers)
                .map(|_| {
                    (
                        BufferCache::new(
                            settings.vertices_buffer_size(),
                            settings.triangles_buffer_size(),
                            buffers,
                        ),
                        BufferCacheAvailable::Available,
                    )
                })
                .collect(),
            _marker: default(),
        }
    }

    fn return_buffer(&mut self, buffer_cache: BufferCache) {
        if self.buffer_cache_availability.contains_key(&buffer_cache) {
            self.buffer_cache_availability
                .insert(buffer_cache, BufferCacheAvailable::Available);
        }
    }
}

fn drain_chunks_to_load<Sampler: Send + Sync>(
    cache: &mut ChunkGeneratorCache<Sampler>,
    gpu_cache: &mut GpuCpuChunkGeneratorCache<Sampler>,
    settings: &ChunkGeneratorSettings<Sampler>,
) -> impl Iterator<Item = (IVec3, BufferCache)> {
    let num_chunks = if matches!(settings.running, ChunkGeneratorRunning::Pause) {
        0
    } else {
        cache
            .chunks_to_load
            .len()
            .min(settings.max_chunks_per_frame)
            .min(
                gpu_cache
                    .buffer_cache_availability
                    .iter()
                    .filter(|(_, a)| matches!(a, BufferCacheAvailable::Available))
                    .count(),
            )
    };

    let buffers = gpu_cache
        .buffer_cache_availability
        .iter()
        .filter(|(_, a)| matches!(a, BufferCacheAvailable::Available))
        .map(|(b, _)| b.clone())
        .take(num_chunks)
        .collect::<Vec<_>>();
    for buffer in buffers.iter() {
        gpu_cache
            .buffer_cache_availability
            .insert(buffer.clone(), BufferCacheAvailable::Unavailable);
    }

    cache.chunks_to_load.drain(..num_chunks).zip(buffers)
}

/// Holds extra buffers to be used when sampling the density field. Gpu Side.
#[derive(Resource, Debug, Clone)]
pub struct GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache> {
    buffer_cache: HashMap<BufferCache, GpuBufferCache>,
    extra_buffer_cache: HashMap<BufferCache, ExtraBufferCache>,
    extra_buffer_needed_keys: Vec<BufferCache>,
    _marker: std::marker::PhantomData<Sampler>,
}

impl<Sampler: Send + Sync, ExtraBufferCache> GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache> {
    pub fn drain_needed_extra_buffers(&mut self) -> Vec<BufferCache> {
        self.extra_buffer_needed_keys.drain(..).collect()
    }
}

impl<Sampler, ExtraBufferCache> Default for GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache> {
    fn default() -> Self {
        Self {
            buffer_cache: default(),
            extra_buffer_cache: default(),
            extra_buffer_needed_keys: default(),
            _marker: default(),
        }
    }
}

impl<Sampler, ExtraBufferCache: GpuExtraBufferCache>
    GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache>
{
    pub fn insert_extra_buffers(&mut self, key: BufferCache, buffers: ExtraBufferCache) {
        self.extra_buffer_cache.insert(key, buffers);
    }

    pub fn extra_buffers_mut(&mut self) -> impl Iterator<Item = &mut ExtraBufferCache> {
        self.extra_buffer_cache.values_mut()
    }
}

#[derive(Debug, Clone)]
struct GpuBufferCache {
    chunk_position_buffer: Buffer,
    sample_pipeline: CachedComputePipelineId,
    sample_workgroups: u32,
    sample_bind_group: BindGroup,
    march_pipeline: CachedComputePipelineId,
    march_workgroups: u32,
    march_bind_group: BindGroup,
}

impl GpuBufferCache {
    fn new<Sampler: Send + Sync + 'static>(
        chunk: &ChunkRenderData<Sampler>,
        extra_buffers: Option<Vec<Buffer>>,
        settings: &ChunkGeneratorSettings<Sampler>,
        render_device: &RenderDevice,
        render_queue: &RenderQueue,
        buffers: &RenderAssets<GpuShaderBuffer>,
        pipelines: &ChunkGeneratorComputePipelines<Sampler>,
        pipeline_cache: &PipelineCache,
    ) -> Self {
        let num_voxels_per_axis = settings.num_voxels_per_axis;
        let num_samples_per_axis = settings.num_samples_per_axis();
        let chunk_size = settings.chunk_size;
        let surface_threshold = settings.surface_threshold;
        let sample_workgroups = (num_samples_per_axis as f32 / WORKGROUP_SIZE as f32).ceil() as u32;
        let march_workgroups = (num_voxels_per_axis as f32 / WORKGROUP_SIZE as f32).ceil() as u32;

        let mut chunk_position_buffer = UniformBuffer::from(chunk.position);
        chunk_position_buffer.write_buffer(render_device, render_queue);

        let mut settings_buffer = UniformBuffer::from(MeshSettings {
            num_voxels_per_axis,
            num_samples_per_axis,
            chunk_size,
            surface_threshold,
        });
        settings_buffer.write_buffer(render_device, render_queue);

        let mut densities_buffer = StorageBuffer::from(vec![
            0.0f32;
            settings.num_samples_per_axis().pow(3)
                as usize
        ]);
        densities_buffer.write_buffer(render_device, render_queue);

        let vertices_buffer = buffers.get(&chunk.buffers.vertices).unwrap();
        let num_vertices_buffer = buffers.get(&chunk.buffers.num_vertices).unwrap();
        let triangles_buffer = buffers.get(&chunk.buffers.triangles).unwrap();
        let num_triangles_buffer = buffers.get(&chunk.buffers.num_triangles).unwrap();

        let sample_bind_group = render_device.create_bind_group(
            Some("marching cubes sample bind group"),
            &pipeline_cache.get_bind_group_layout(&pipelines.sample_layout),
            &[
                chunk_position_buffer.binding().unwrap(),
                settings_buffer.binding().unwrap(),
                densities_buffer.binding().unwrap(),
            ]
            .into_iter()
            .chain(
                extra_buffers
                    .unwrap_or(vec![])
                    .iter()
                    .map(|b| b.as_entire_binding()),
            )
            .enumerate()
            .map(|(i, res)| BindGroupEntry {
                binding: i as u32,
                resource: res,
            })
            .collect::<Vec<_>>(),
        );

        let march_bind_group = render_device.create_bind_group(
            Some("marching cubes march bind group"),
            &pipeline_cache.get_bind_group_layout(&pipelines.march_layout),
            &BindGroupEntries::sequential((
                densities_buffer.binding().unwrap(),
                settings_buffer.binding().unwrap(),
                vertices_buffer.buffer.as_entire_binding(),
                num_vertices_buffer.buffer.as_entire_binding(),
                triangles_buffer.buffer.as_entire_binding(),
                num_triangles_buffer.buffer.as_entire_binding(),
            )),
        );

        GpuBufferCache {
            chunk_position_buffer: chunk_position_buffer.buffer().unwrap().clone(),
            sample_pipeline: pipelines.sample_pipeline,
            sample_workgroups,
            sample_bind_group,
            march_pipeline: pipelines.march_pipeline,
            march_workgroups,
            march_bind_group,
        }
    }
}

/// Event meaning any extra buffers should be cleared entirely to prevent memory leaks, like when the chunk generator is destroyed.
#[derive(Event)]
pub struct ClearBufferCache<Sampler>(std::marker::PhantomData<Sampler>);

fn init_cache<Sampler: Send + Sync + 'static>(
    mut commands: Commands,
    mut settings: ResMut<ChunkGeneratorSettings<Sampler>>,
    mut gpu_settings: ResMut<GpuChunkGeneratorSettings<Sampler>>,
    mut buffers: ResMut<Assets<ShaderBuffer>>,
    cache: Option<Res<ChunkGeneratorCache<Sampler>>>,
) {
    let mut needs_cache = cache.is_none();
    if matches!(
        settings.running,
        ChunkGeneratorRunning::Stop | ChunkGeneratorRunning::Reset
    ) {
        gpu_settings.clear_gpu_cache = true;
        commands.trigger(ClearBufferCache::<Sampler>(default()));
        commands.remove_resource::<ChunkGeneratorCache<Sampler>>();
        commands.remove_resource::<GpuCpuChunkGeneratorCache<Sampler>>();
        needs_cache = true;
    } else {
        gpu_settings.clear_gpu_cache = false;
    }

    if matches!(settings.running, ChunkGeneratorRunning::Reset) {
        settings.running = ChunkGeneratorRunning::Run;
    }

    if matches!(
        settings.running,
        ChunkGeneratorRunning::Run | ChunkGeneratorRunning::Pause
    ) && needs_cache
    {
        commands.insert_resource(ChunkGeneratorCache::<Sampler>::default());
        commands.insert_resource(GpuCpuChunkGeneratorCache::<Sampler>::new(
            &settings,
            &mut buffers,
        ));
    }
}

fn clear_old_dispatches(mut dispatches: ResMut<ChunkGeneratorDispatches>) {
    dispatches.dispatches.clear();
}

fn check_pipelines<Sampler: Send + Sync + 'static>(
    pipelines: Option<Res<ChunkGeneratorComputePipelines<Sampler>>>,
    pipeline_cache: Res<PipelineCache>,
    mut pipelines_done_loading: ResMut<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>,
) {
    pipelines_done_loading.done = pipelines.is_some_and(|pipelines| {
        pipeline_cache
            .get_compute_pipeline(pipelines.sample_pipeline)
            .is_some()
            && pipeline_cache
                .get_compute_pipeline(pipelines.march_pipeline)
                .is_some()
    });
}

fn calc_needed_buffers<
    Sampler: GpuChunkComputer + Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
>(
    chunks: Query<(Entity, &ChunkRenderData<Sampler>)>,
    settings: Res<GpuChunkGeneratorSettings<Sampler>>,
    mut cache: ResMut<GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache>>,
    mut processed: Local<HashSet<IVec3>>,
    mut commands: Commands,
) -> Result<()> {
    if settings.clear_gpu_cache {
        processed.clear();
    }

    for (entity, chunk) in chunks.iter() {
        if processed.contains(&chunk.position) {
            continue;
        }
        processed.insert(chunk.position);

        trace!(
            "calc_needed_buffers {} {}",
            ShortName::of::<Sampler>(),
            chunk.position
        );

        commands.entity(entity).insert(ChunkThisFrame);

        if !cache.extra_buffer_cache.contains_key(&chunk.buffers) {
            cache.extra_buffer_needed_keys.push(chunk.buffers.clone());
        }
    }

    Ok(())
}

fn prepare_bind_groups<
    Sampler: GpuChunkComputer + Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
>(
    render_device: Res<RenderDevice>,
    render_queue: Res<RenderQueue>,
    chunks: Query<&ChunkRenderData<Sampler>, With<ChunkThisFrame>>,
    buffers: Res<RenderAssets<GpuShaderBuffer>>,
    settings: Res<ChunkGeneratorSettings<Sampler>>,
    pipelines: Res<ChunkGeneratorComputePipelines<Sampler>>,
    mut dispatches: ResMut<ChunkGeneratorDispatches>,
    mut cache: ResMut<GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache>>,
    pipeline_cache: Res<PipelineCache>,
) -> Result<()> {
    for chunk in chunks.iter() {
        let extra_buffers = cache
            .extra_buffer_cache
            .get(&chunk.buffers)
            .map(|b| b.buffers());

        trace!(
            "prepare_bind_groups {} {} with {:?} extra buffers, currently {} buffer sets loaded",
            ShortName::of::<Sampler>(),
            chunk.position,
            extra_buffers.as_ref().map(|b| b.len()),
            cache.buffer_cache.len()
        );

        let buffer_cache = cache
            .buffer_cache
            .entry(chunk.buffers.clone())
            .or_insert_with(|| {
                GpuBufferCache::new(
                    chunk,
                    extra_buffers,
                    &settings,
                    &render_device,
                    &render_queue,
                    &buffers,
                    &pipelines,
                    &pipeline_cache,
                )
            })
            .clone();

        render_queue.write_buffer(
            &buffer_cache.chunk_position_buffer,
            0,
            &value_data(&chunk.position)?,
        );
        let num_vertices_buffer = buffers.get(&chunk.buffers.num_vertices).unwrap();
        render_queue.write_buffer(&num_vertices_buffer.buffer, 0, &value_data(&0u32)?);
        let num_triangles_buffer = buffers.get(&chunk.buffers.num_triangles).unwrap();
        render_queue.write_buffer(&num_triangles_buffer.buffer, 0, &value_data(&0u32)?);

        dispatches.dispatches.push(buffer_cache);
    }

    Ok(())
}

/// Converts a [`ShaderType`] into [`Vec<u8>`] data.
pub fn value_data<T: ShaderType + WriteInto>(value: &T) -> Result<Vec<u8>> {
    let mut writer = encase::StorageBuffer::<Vec<u8>>::new(Vec::new());
    writer.write(value)?;
    Ok(writer.into_inner())
}

/// Event that needs to be sent when an extra buffer has been finished reading from.
/// Otherwise, the chunk will never finish generating.
#[derive(EntityEvent)]
pub struct ReadbackReallyComplete(pub Entity);

fn clear_gpu_cache<
    Sampler: GpuChunkComputer + Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
>(
    settings: Res<GpuChunkGeneratorSettings<Sampler>>,
    mut cache: ResMut<GpuGpuChunkGeneratorCache<Sampler, ExtraBufferCache>>,
) {
    if settings.clear_gpu_cache {
        cache.buffer_cache.clear();
        cache.extra_buffer_cache.clear();
        cache.extra_buffer_needed_keys.clear();
    }
}

#[derive(Resource, Debug)]
struct ChunkGeneratorComputePipelinesDoneLoading<Sampler> {
    done: bool,
    _marker: std::marker::PhantomData<Sampler>,
}

impl<Sampler> Default for ChunkGeneratorComputePipelinesDoneLoading<Sampler> {
    fn default() -> Self {
        Self {
            done: false,
            _marker: default(),
        }
    }
}

impl<Sampler> Clone for ChunkGeneratorComputePipelinesDoneLoading<Sampler> {
    fn clone(&self) -> Self {
        Self {
            done: self.done,
            _marker: self._marker,
        }
    }
}

fn extract_pipelines_done<Sampler: Send + Sync + 'static>(
    render_resource: Option<Res<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>>,
    mut main_world: ResMut<MainWorld>,
) {
    if let Some(render_resource) = render_resource.as_ref() {
        if let Some(mut target_resource) =
            main_world.get_resource_mut::<ChunkGeneratorComputePipelinesDoneLoading<Sampler>>()
        {
            if render_resource.is_changed() {
                *target_resource = (*render_resource).clone();
            }
        } else {
            main_world.insert_resource((*render_resource).clone());
        }
    }
}

#[derive(ExtractResource, Resource, Debug)]
struct GpuChunkGeneratorSettings<Sampler: Send + Sync + 'static> {
    clear_gpu_cache: bool,
    _marker: std::marker::PhantomData<Sampler>,
}

impl<Sampler: Send + Sync + 'static> Clone for GpuChunkGeneratorSettings<Sampler> {
    fn clone(&self) -> Self {
        Self {
            clear_gpu_cache: self.clear_gpu_cache,
            _marker: self._marker,
        }
    }
}

impl<Sampler: Send + Sync + 'static> Default for GpuChunkGeneratorSettings<Sampler> {
    fn default() -> Self {
        Self {
            clear_gpu_cache: false,
            _marker: PhantomData,
        }
    }
}

fn start_chunks<
    Sampler: Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
    Material: Asset + bevy::prelude::Material,
>(
    mut commands: Commands,
    settings: Res<ChunkGeneratorSettings<Sampler>>,
    mut cache: ResMut<ChunkGeneratorCache<Sampler>>,
    mut gpu_cache: ResMut<GpuCpuChunkGeneratorCache<Sampler>>,
    material: Res<ChunkMaterial<Sampler, Material>>,
) {
    for (chunk_position, buffers) in drain_chunks_to_load(&mut cache, &mut gpu_cache, &settings) {
        trace!("start_chunks {chunk_position:?}");

        let chunk_entity = commands
            .spawn((
                Name::new(format!("Chunk {chunk_position:?}")),
                Transform::from_translation(settings.chunk_to_position(chunk_position)),
                MeshMaterial3d(material.material.clone()),
                Chunk::<Sampler> {
                    position: chunk_position,
                    _marker: default(),
                },
                ChunkGenData {
                    vertices: None,
                    triangles: None,
                },
                ChunkRenderData::<Sampler> {
                    position: chunk_position,
                    buffers: buffers.clone(),
                    num_readbacks_done: 0,
                    _marker: default(),
                },
            ))
            .observe(finish_chunk::<Sampler, ExtraBufferCache>)
            .id();

        commands
            .spawn((
                Name::new(format!("Chunk {chunk_position:?} num_vertices readback")),
                Readback::buffer(buffers.num_vertices),
                ChildOf(chunk_entity),
            ))
            .observe(
                move |readback: On<ReadbackComplete>,
                      mut chunks: Query<&mut ChunkGenData>,
                      mut commands: Commands|
                      -> Result {
                    let mut chunk = chunks.get_mut(chunk_entity)?;
                    let num_vertices: u32 = readback.to_shader_type();
                    trace!("num_vertices readback {chunk_position:?} {num_vertices}");
                    commands.entity(readback.entity).despawn();
                    if num_vertices > 0 {
                        commands
                            .spawn((
                                Name::new(format!("Chunk {chunk_position:?} vertices readback")),
                                Readback::buffer_range(
                                    buffers.vertices.clone(),
                                    0,
                                    size_of::<Vertex>() as u64 * num_vertices as u64,
                                ),
                                ChildOf(chunk_entity),
                            ))
                            .observe(
                                move |readback: On<ReadbackComplete>,
                                      mut chunks: Query<&mut ChunkGenData>,
                                      mut commands: Commands|
                                      -> Result {
                                    let vertices: Vec<Vertex> = readback.to_shader_type();
                                    trace!(
                                        "vertices readback {chunk_position:?} {}",
                                        vertices.len()
                                    );
                                    let mut chunk = chunks.get_mut(chunk_entity)?;
                                    chunk.vertices = Some(vertices);
                                    commands.trigger(ReadbackReallyComplete(chunk_entity));
                                    commands.entity(readback.entity).despawn();
                                    Ok(())
                                },
                            );
                    } else {
                        chunk.vertices = Some(vec![]);
                        commands.trigger(ReadbackReallyComplete(chunk_entity));
                    }
                    Ok(())
                },
            );

        commands
            .spawn((
                Name::new(format!("Chunk {chunk_position:?} num_triangles readback")),
                Readback::buffer(buffers.num_triangles),
                ChildOf(chunk_entity),
            ))
            .observe(
                move |readback: On<ReadbackComplete>,
                      mut chunks: Query<&mut ChunkGenData>,
                      mut commands: Commands|
                      -> Result {
                    let mut chunk = chunks.get_mut(chunk_entity)?;
                    let num_triangles: u32 = readback.to_shader_type();
                    trace!("num_triangles readback {chunk_position:?} {num_triangles}");
                    commands.entity(readback.entity).despawn();
                    if num_triangles > 0 {
                        commands
                            .spawn((
                                Name::new(format!("Chunk {chunk_position:?} triangles readback")),
                                Readback::buffer_range(
                                    buffers.triangles.clone(),
                                    0,
                                    size_of::<Triangle>() as u64 * num_triangles as u64,
                                ),
                                ChildOf(chunk_entity),
                            ))
                            .observe(
                                move |readback: On<ReadbackComplete>,
                                      mut chunks: Query<&mut ChunkGenData>,
                                      mut commands: Commands|
                                      -> Result {
                                    let triangles: Vec<Triangle> = readback.to_shader_type();
                                    trace!(
                                        "triangles readback {chunk_position:?} {}",
                                        triangles.len()
                                    );
                                    let mut chunk = chunks.get_mut(chunk_entity)?;
                                    chunk.triangles = Some(triangles);
                                    commands.trigger(ReadbackReallyComplete(chunk_entity));
                                    commands.entity(readback.entity).despawn();
                                    Ok(())
                                },
                            );
                    } else {
                        chunk.triangles = Some(vec![]);
                        commands.trigger(ReadbackReallyComplete(chunk_entity));
                    }
                    Ok(())
                },
            );
    }
}

fn finish_chunk<
    Sampler: Send + Sync + 'static,
    ExtraBufferCache: GpuExtraBufferCache + Send + Sync + 'static,
>(
    readback: On<ReadbackReallyComplete>,
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut cache: ResMut<ChunkGeneratorCache<Sampler>>,
    mut gpu_cache: ResMut<GpuCpuChunkGeneratorCache<Sampler>>,
    mut chunks: Query<(&mut ChunkRenderData<Sampler>, &ChunkGenData)>,
) -> Result {
    let chunk_entity = readback.0;
    let Ok((mut chunk, chunk_gen)) = chunks.get_mut(chunk_entity) else {
        return Ok(()); // It's already done
    };

    chunk.num_readbacks_done += 1;

    if chunk.num_readbacks_done < ExtraBufferCache::num_extra_readbacks() + 2 {
        return Ok(());
    }
    let Some(ref vertices) = chunk_gen.vertices else {
        return Ok(());
    };
    let Some(ref triangles) = chunk_gen.triangles else {
        return Ok(());
    };

    trace!("finish_chunk {}", chunk.position);

    if !vertices.is_empty() && !triangles.is_empty() {
        let mesh = Mesh::new(
            bevy::mesh::PrimitiveTopology::TriangleList,
            bevy::asset::RenderAssetUsages::RENDER_WORLD,
        )
        .with_inserted_indices(bevy::mesh::Indices::U32(
            triangles
                .iter()
                .flat_map(|t| [t.vertex_c, t.vertex_b, t.vertex_a])
                .collect(),
        ))
        .with_inserted_attribute(
            Mesh::ATTRIBUTE_POSITION,
            vertices.iter().map(|v| v.position).collect::<Vec<_>>(),
        )
        .with_inserted_attribute(
            Mesh::ATTRIBUTE_NORMAL,
            vertices.iter().map(|v| v.normal).collect::<Vec<_>>(),
        )
        .with_inserted_attribute(
            Mesh::ATTRIBUTE_UV_0,
            vertices.iter().map(|v| v.position.xy()).collect::<Vec<_>>(),
        );

        commands
            .entity(chunk_entity)
            .insert(Mesh3d(meshes.add(mesh)));
    }
    cache
        .loaded_chunks
        .insert(chunk.position, LoadState::Finished);

    gpu_cache.return_buffer(chunk.buffers.clone());
    commands
        .entity(chunk_entity)
        .remove::<ChunkGenData>()
        .remove::<ChunkRenderData<Sampler>>();
    commands.trigger(ChunkGenerated(chunk_entity));

    Ok(())
}

#[derive(Component, Debug)]
struct ChunkGenData {
    vertices: Option<Vec<Vertex>>,
    triangles: Option<Vec<Triangle>>,
}

#[derive(encase::ShaderType)]
struct MeshSettings {
    num_voxels_per_axis: u32,
    num_samples_per_axis: u32,
    chunk_size: f32,
    surface_threshold: f32,
}

#[derive(Component, Debug, Default)]
struct ChunkThisFrame;