bevy_cube_marcher 0.5.1

use bevy::{
    ecs::{component::Mutable, world::CommandQueue},
    tasks::{AsyncComputeTaskPool, Task, futures::check_ready},
};

use crate::{cpu::marching_cubes::marching_cubes, *};

mod marching_cubes;

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

impl Plugin for CpuMarchingCubesGlobalPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Update, handle_tasks.in_set(ChunkGenSystems));
    }
}

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

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

impl<
    Sampler: Resource<Mutability = Mutable> + CpuChunkComputer + Clone + Default,
    Material: Asset + bevy::prelude::Material,
> Plugin for CpuMarchingCubesPlugin<Sampler, Material>
{
    fn build(&self, app: &mut App) {
        if !app.is_plugin_added::<CpuMarchingCubesGlobalPlugin>() {
            app.add_plugins(CpuMarchingCubesGlobalPlugin);
        }

        app.init_resource::<Sampler>().add_systems(
            Update,
            (
                reset_generator::<Sampler>,
                (
                    update_chunk_loaders::<Sampler>,
                    queue_chunks::<Sampler>,
                    start_chunks::<Sampler, Material>,
                )
                    .chain()
                    .run_if(resource_exists::<ChunkGeneratorCache<Sampler>>),
            )
                .chain()
                .in_set(ChunkGenSystems),
        );
    }
}

fn reset_generator<Sampler: Send + Sync + 'static>(
    mut commands: Commands,
    mut settings: ResMut<ChunkGeneratorSettings<Sampler>>,
    cache: Option<Res<ChunkGeneratorCache<Sampler>>>,
) {
    let mut needs_cache = cache.is_none();
    if matches!(
        settings.running,
        ChunkGeneratorRunning::Stop | ChunkGeneratorRunning::Reset
    ) {
        commands.remove_resource::<ChunkGeneratorCache<Sampler>>();
        needs_cache = true;
    }

    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());
    }
}

fn drain_chunks_to_load<Sampler: Send + Sync>(
    cache: &mut ChunkGeneratorCache<Sampler>,
    settings: &ChunkGeneratorSettings<Sampler>,
) -> impl Iterator<Item = IVec3> {
    let num_chunks = if matches!(settings.running, ChunkGeneratorRunning::Pause) {
        0
    } else {
        cache
            .chunks_to_load
            .len()
            .min(settings.max_chunks_per_frame)
    };

    cache.chunks_to_load.drain(..num_chunks)
}

fn start_chunks<
    Sampler: Resource<Mutability = Mutable> + CpuChunkComputer + Clone,
    Material: Asset + bevy::prelude::Material,
>(
    mut commands: Commands,
    settings: Res<ChunkGeneratorSettings<Sampler>>,
    mut cache: ResMut<ChunkGeneratorCache<Sampler>>,
    material: Res<ChunkMaterial<Sampler, Material>>,
    sampler: ResMut<Sampler>,
) {
    let task_pool = AsyncComputeTaskPool::get();
    for chunk_position in drain_chunks_to_load(&mut 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(),
                },
            ))
            .id();

        let settings = settings.clone();
        let sampler = sampler.clone();
        commands
            .entity(chunk_entity)
            .insert(ComputeChunk(task_pool.spawn(async move {
                let mut command_queue = CommandQueue::default();

                let num_samples_per_axis = settings.num_samples_per_axis() as usize;
                let num_samples =
                    num_samples_per_axis * num_samples_per_axis * num_samples_per_axis;
                let mut densities = vec![0.0; num_samples];
                for x in 0..num_samples_per_axis {
                    for y in 0..num_samples_per_axis {
                        for z in 0..num_samples_per_axis {
                            densities[x * num_samples_per_axis * num_samples_per_axis
                                + y * num_samples_per_axis
                                + z] = sampler.sample(
                                UVec3::new(x as u32, y as u32, z as u32),
                                chunk_position,
                                &settings,
                            );
                        }
                    }
                }

                let num_voxels_per_axis = settings.num_voxels_per_axis;
                let mut vertices = vec![];
                let mut triangles = vec![];
                for x in 0..num_voxels_per_axis {
                    for y in 0..num_voxels_per_axis {
                        for z in 0..num_voxels_per_axis {
                            marching_cubes(
                                UVec3::new(x, y, z),
                                &densities,
                                &mut vertices,
                                &mut triangles,
                                &settings,
                            );
                        }
                    }
                }

                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<_>>(),
                    );

                    command_queue.push(move |world: &mut World| {
                        let mut meshes: Mut<Assets<Mesh>> = world.resource_mut();
                        let mesh = meshes.add(mesh);
                        world.entity_mut(chunk_entity).insert(Mesh3d(mesh));
                    });
                }

                command_queue.push(move |world: &mut World| {
                    let mut cache: Mut<ChunkGeneratorCache<Sampler>> = world.resource_mut();
                    cache
                        .loaded_chunks
                        .insert(chunk_position, LoadState::Finished);
                    world.trigger(ChunkGenerated(chunk_entity));
                });

                command_queue
            })));
    }
}

#[derive(Component)]
struct ComputeChunk(Task<CommandQueue>);

fn handle_tasks(mut commands: Commands, tasks: Query<(Entity, &mut ComputeChunk)>) {
    for (entity, mut task) in tasks {
        if let Some(mut command_queue) = check_ready(&mut task.0) {
            commands.append(&mut command_queue);
            commands.entity(entity).remove::<ComputeChunk>();
        }
    }
}

/// Samples the density field.
pub trait CpuChunkComputer: Send + Sync + Sized {
    fn sample(
        &self,
        coord: UVec3,
        chunk_position: IVec3,
        settings: &ChunkGeneratorSettings<Self>,
    ) -> f32;
}