//! Demonstrates how to define and use specialized mesh pipeline
//!
//! This example shows how to use the built-in [`SpecializedMeshPipeline`]
//! functionality with a custom [`RenderCommand`](bevy::render::render_phase::RenderCommand)
//! to allow custom mesh rendering with more flexibility than the material api.
//!
//! [`SpecializedMeshPipeline`] let's you customize the entire pipeline used when rendering a mesh.

use bevy::{
    asset::RenderAssetUsages,
    camera::visibility::{self, VisibilityClass},
    core_pipeline::core_3d::{Opaque3d, Opaque3dBatchSetKey, Opaque3dBinKey, CORE_3D_DEPTH_FORMAT},
    ecs::change_detection::Tick,
    math::{vec3, vec4},
    mesh::{
        Indices, MeshAttributeCompressionFlags, MeshCompressionArgs, MeshVertexBufferLayoutRef,
        PrimitiveTopology,
    },
    pbr::{
        DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineSystems, MeshPipelineViewLayoutKey,
        RenderMeshInstances, SetMeshBindGroup, SetMeshViewBindGroup, SetMeshViewEmptyBindGroup,
        ViewKeyCache,
    },
    prelude::*,
    render::{
        batching::gpu_preprocessing::GpuPreprocessingSupport,
        camera::{DirtySpecializations, PendingQueues},
        extract_component::{ExtractComponent, ExtractComponentPlugin},
        mesh::{allocator::MeshAllocator, RenderMesh},
        render_asset::RenderAssets,
        render_phase::{
            AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, SetItemPipeline,
            ViewBinnedRenderPhases,
        },
        render_resource::{
            ColorTargetState, ColorWrites, CompareFunction, DepthStencilState, Face, FragmentState,
            FrontFace, MultisampleState, PipelineCache, PolygonMode, PrimitiveState,
            RenderPipelineDescriptor, SpecializedMeshPipeline, SpecializedMeshPipelineError,
            SpecializedMeshPipelines, VertexState,
        },
        sync_world::MainEntityHashSet,
        view::{ExtractedView, RenderVisibleEntities},
        Render, RenderApp, RenderStartup, RenderSystems,
    },
};

const SHADER_ASSET_PATH: &str = "shaders/specialized_mesh_pipeline.wesl";

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugins(CustomRenderedMeshPipelinePlugin)
        .add_systems(Startup, setup)
        .run();
}

/// Spawns the objects in the scene.
fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
    // Build a custom triangle mesh with colors
    // We define a custom mesh because the examples only uses a limited
    // set of vertex attributes for simplicity
    let mesh = Mesh::new(
        PrimitiveTopology::TriangleList,
        RenderAssetUsages::default(),
    )
    .with_inserted_indices(Indices::U32(vec![0, 1, 2]))
    .with_inserted_attribute(
        Mesh::ATTRIBUTE_POSITION,
        vec![
            vec3(-0.5, -0.5, 0.0),
            vec3(0.5, -0.5, 0.0),
            vec3(0.0, 0.25, 0.0),
        ],
    )
    .with_inserted_attribute(
        Mesh::ATTRIBUTE_COLOR,
        vec![
            vec4(1.0, 0.0, 0.0, 1.0),
            vec4(0.0, 1.0, 0.0, 1.0),
            vec4(0.0, 0.0, 1.0, 1.0),
        ],
    );

    // spawn 3 triangles to show that batching works
    for (x, y) in [-0.5, 0.0, 0.5].into_iter().zip([-0.25, 0.5, -0.25]) {
        // Spawn an entity with all the required components for it to be rendered with our custom pipeline
        commands.spawn((
            // We use a marker component to identify the mesh that will be rendered
            // with our specialized pipeline
            CustomRenderedEntity,
            // We need to add the mesh handle to the entity
            Mesh3d(meshes.add(mesh.clone())),
            Transform::from_xyz(x, y, 0.0),
        ));
    }
    commands.spawn((
        // We use a marker component to identify the mesh that will be rendered
        // with our specialized pipeline
        CustomRenderedEntity,
        // We need to add the mesh handle to the entity
        Mesh3d(
            meshes.add(
                Rectangle::new(0.5, 0.5)
                    .mesh()
                    .build()
                    .with_removed_attribute(Mesh::ATTRIBUTE_NORMAL)
                    .with_removed_attribute(Mesh::ATTRIBUTE_UV_0)
                    .compressed_mesh(&MeshCompressionArgs::regular())
                    .unwrap(),
            ),
        ),
        Transform::from_xyz(0.0, -1.0 / 3.0, -1.0),
    ));

    // Spawn the camera.
    commands.spawn((
        Camera3d::default(),
        // Move the camera back a bit to see all the triangles
        Transform::from_xyz(0.0, 0.0, 3.0).looking_at(Vec3::ZERO, Vec3::Y),
    ));
}

// When writing custom rendering code it's generally recommended to use a plugin.
// The main reason for this is that it gives you access to the finish() hook
// which is called after rendering resources are initialized.
struct CustomRenderedMeshPipelinePlugin;
impl Plugin for CustomRenderedMeshPipelinePlugin {
    fn build(&self, app: &mut App) {
        app.add_plugins(ExtractComponentPlugin::<CustomRenderedEntity>::default());

        // We make sure to add these to the render app, not the main app.
        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };
        render_app
            // This is needed to tell bevy about your custom pipeline
            .init_resource::<SpecializedMeshPipelines<CustomMeshPipeline>>()
            .init_resource::<PendingCustomMeshQueues>()
            // We need to use a custom draw command so we need to register it
            .add_render_command::<Opaque3d, DrawSpecializedPipelineCommands>()
            .add_systems(
                RenderStartup,
                init_custom_mesh_pipeline.after(MeshPipelineSystems),
            )
            .add_systems(
                Render,
                queue_custom_mesh_pipeline.in_set(RenderSystems::Queue),
            );
    }
}

/// A marker component that represents an entity that is to be rendered using
/// our specialized pipeline.
///
/// Note the [`ExtractComponent`] trait implementation: this is necessary to
/// tell Bevy that this object should be pulled into the render world. Also note
/// the `on_add` hook, which is needed to tell Bevy's `check_visibility` system
/// that entities with this component need to be examined for visibility.
#[derive(Clone, Component, ExtractComponent)]
#[require(VisibilityClass)]
#[component(on_add = visibility::add_visibility_class::<CustomRenderedEntity>)]
#[extract_app(RenderApp)]
struct CustomRenderedEntity;

/// The custom draw commands that Bevy executes for each entity we enqueue into
/// the render phase.
type DrawSpecializedPipelineCommands = (
    // Set the pipeline
    SetItemPipeline,
    // Set the view uniform at bind group 0
    SetMeshViewBindGroup<0>,
    // Set an empty material bind group at bind group 1
    SetMeshViewEmptyBindGroup<1>,
    // Set the mesh uniform at bind group 2
    SetMeshBindGroup<2>,
    // Draw the mesh
    DrawMesh,
);

// This contains the state needed to specialize a mesh pipeline
#[derive(Resource)]
struct CustomMeshPipeline {
    /// The base mesh pipeline defined by bevy
    ///
    /// This isn't required, but if you want to use a bevy `Mesh` it's easier when you
    /// have access to the base `MeshPipeline` that bevy already defines
    mesh_pipeline: MeshPipeline,
    /// Stores the shader used for this pipeline directly on the pipeline.
    /// This isn't required, it's only done like this for simplicity.
    shader_handle: Handle<Shader>,
}

fn init_custom_mesh_pipeline(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    mesh_pipeline: Res<MeshPipeline>,
) {
    // Load the shader
    let shader_handle: Handle<Shader> = asset_server.load(SHADER_ASSET_PATH);
    commands.insert_resource(CustomMeshPipeline {
        mesh_pipeline: mesh_pipeline.clone(),
        shader_handle,
    });
}

impl SpecializedMeshPipeline for CustomMeshPipeline {
    /// Pipeline use keys to determine how to specialize it.
    /// The key is also used by the pipeline cache to determine if
    /// it needs to create a new pipeline or not
    ///
    /// In this example we just use the base `MeshPipelineKey` defined by bevy, but this could be anything.
    /// For example, if you want to make a pipeline with a procedural shader you could add the Handle<Shader> to the key.
    type Key = MeshPipelineKey;

    fn specialize(
        &self,
        mesh_key: Self::Key,
        layout: &MeshVertexBufferLayoutRef,
    ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
        let mut shader_defs = Vec::new();
        // Define the vertex attributes based on a standard bevy [`Mesh`]
        let mut vertex_attributes = Vec::new();
        if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
            // Handle compressed vertex positions.
            if layout
                .0
                .get_attribute_compression()
                .contains(MeshAttributeCompressionFlags::COMPRESS_POSITION)
            {
                shader_defs.push("VERTEX_POSITIONS_COMPRESSED".into());
            }
            // Make sure this matches the shader location
            vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
        }
        if layout.0.contains(Mesh::ATTRIBUTE_COLOR) {
            shader_defs.push("VERTEX_COLORS".into());
            // Make sure this matches the shader location
            vertex_attributes.push(Mesh::ATTRIBUTE_COLOR.at_shader_location(1));
        }
        // This will automatically generate the correct `VertexBufferLayout` based on the vertex attributes
        let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;

        let view_layout = self
            .mesh_pipeline
            .get_view_layout(MeshPipelineViewLayoutKey::from(mesh_key));

        Ok(RenderPipelineDescriptor {
            label: Some("Specialized Mesh Pipeline".into()),
            layout: vec![
                view_layout.main_layout,
                view_layout.empty_layout,
                self.mesh_pipeline.mesh_layouts.model_only.clone(),
            ],
            vertex: VertexState {
                shader: self.shader_handle.clone(),
                shader_defs: shader_defs.clone(),
                // Customize how to store the meshes' vertex attributes in the vertex buffer
                buffers: vec![vertex_buffer_layout],
                ..default()
            },
            fragment: Some(FragmentState {
                shader: self.shader_handle.clone(),
                shader_defs,
                targets: vec![Some(ColorTargetState {
                    // This isn't required, but bevy supports rendering different formats
                    // so it's generally recommended to specialize the pipeline for that
                    format: mesh_key.target_format(),
                    // For this example we only use opaque meshes,
                    // but if you wanted to use alpha blending you would need to set it here
                    blend: None,
                    write_mask: ColorWrites::ALL,
                })],
                ..default()
            }),
            primitive: PrimitiveState {
                topology: mesh_key.primitive_topology(),
                strip_index_format: mesh_key.strip_index_format(),
                front_face: FrontFace::Ccw,
                cull_mode: Some(Face::Back),
                polygon_mode: PolygonMode::Fill,
                ..default()
            },
            // Note that if your view has no depth buffer this will need to be
            // changed.
            depth_stencil: Some(DepthStencilState {
                format: CORE_3D_DEPTH_FORMAT,
                depth_write_enabled: Some(true),
                depth_compare: Some(CompareFunction::GreaterEqual),
                stencil: default(),
                bias: default(),
            }),
            // It's generally recommended to specialize your pipeline for MSAA,
            // but it's not always possible
            multisample: MultisampleState {
                count: mesh_key.msaa_samples(),
                ..default()
            },
            ..default()
        })
    }
}

/// A resource that stores meshes that couldn't be specialized yet because their
/// materials hadn't loaded.
///
/// See the documentation for [`PendingQueues`] for more information.
#[derive(Default, Deref, DerefMut, Resource)]
struct PendingCustomMeshQueues(pub PendingQueues);

/// A render-world system that enqueues the entity with custom rendering into
/// the opaque render phases of each view.
fn queue_custom_mesh_pipeline(
    pipeline_cache: Res<PipelineCache>,
    custom_mesh_pipeline: Res<CustomMeshPipeline>,
    (mut opaque_render_phases, opaque_draw_functions): (
        ResMut<ViewBinnedRenderPhases<Opaque3d>>,
        Res<DrawFunctions<Opaque3d>>,
    ),
    mut specialized_mesh_pipelines: ResMut<SpecializedMeshPipelines<CustomMeshPipeline>>,
    views: Query<(&RenderVisibleEntities, &ExtractedView)>,
    view_key_cache: Res<ViewKeyCache>,
    (render_meshes, render_mesh_instances): (
        Res<RenderAssets<RenderMesh>>,
        Res<RenderMeshInstances>,
    ),
    mut change_tick: Local<Tick>,
    mesh_allocator: Res<MeshAllocator>,
    gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
    dirty_specializations: Res<DirtySpecializations>,
    mut pending_custom_mesh_queues: ResMut<PendingCustomMeshQueues>,
    mut mesh_instances_queued_this_iteration_scratch_space: Local<MainEntityHashSet>,
) {
    // Get the id for our custom draw function
    let draw_function = opaque_draw_functions
        .read()
        .id::<DrawSpecializedPipelineCommands>();

    // Render phases are per-view, so we need to iterate over all views so that
    // the entity appears in them. (In this example, we have only one view, but
    // it's good practice to loop over all views anyway.)
    for (view_visible_entities, view) in views.iter() {
        let Some(opaque_phase) = opaque_render_phases.get_mut(&view.retained_view_entity) else {
            continue;
        };

        let Some(&view_key) = view_key_cache.get(&view.retained_view_entity) else {
            continue;
        };

        let Some(render_visible_mesh_entities) =
            view_visible_entities.get::<CustomRenderedEntity>()
        else {
            continue;
        };

        // Initialize the pending queues.
        let view_pending_custom_mesh_queues =
            pending_custom_mesh_queues.prepare_for_new_frame(view.retained_view_entity);

        // First, remove meshes that need to be respecialized, and those that were removed, from the bins.
        for &main_entity in dirty_specializations
            .iter_to_dequeue(view.retained_view_entity, render_visible_mesh_entities)
        {
            opaque_phase.remove(main_entity);
        }

        // Find all the custom rendered entities that are visible from this
        // view.
        for (render_entity, visible_entity) in dirty_specializations.iter_to_queue(
            view.retained_view_entity,
            render_visible_mesh_entities,
            &view_pending_custom_mesh_queues.prev_frame,
            &mut mesh_instances_queued_this_iteration_scratch_space,
        ) {
            // Get the mesh instance
            let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
            else {
                // We couldn't fetch the mesh, probably because it hasn't been
                // loaded yet. Add the entity to the list of pending custom
                // meshes and bail.
                view_pending_custom_mesh_queues
                    .current_frame
                    .insert((*render_entity, *visible_entity));
                continue;
            };

            // Get the mesh data
            let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id()) else {
                continue;
            };

            let Some(mesh_slabs) = mesh_allocator.mesh_slabs(&mesh_instance.mesh_asset_id()) else {
                continue;
            };

            // Specialize the key for the current mesh entity
            // For this example we only specialize based on the mesh topology
            // but you could have more complex keys and that's where you'd need to create those keys
            let mut mesh_key = view_key;
            mesh_key |= MeshPipelineKey::from_primitive_topology_and_strip_index(
                mesh.primitive_topology(),
                mesh.index_format(),
            );

            // Finally, we can specialize the pipeline based on the key
            let pipeline_id = specialized_mesh_pipelines
                .specialize(
                    &pipeline_cache,
                    &custom_mesh_pipeline,
                    mesh_key,
                    &mesh.layout,
                )
                // This should never happen with this example, but if your pipeline
                // specialization can fail you need to handle the error here
                .expect("Failed to specialize mesh pipeline");

            // Bump the change tick so that Bevy is forced to rebuild the bin.
            let next_change_tick = change_tick.get() + 1;
            change_tick.set(next_change_tick);

            // Add the mesh with our specialized pipeline
            opaque_phase.add(
                Opaque3dBatchSetKey {
                    draw_function,
                    pipeline: pipeline_id,
                    material_bind_group_index: None,
                    slabs: mesh_slabs,
                    lightmap_slab: None,
                },
                // For this example we can use the mesh asset id as the bin key,
                // but you can use any asset_id as a key
                Opaque3dBinKey {
                    asset_id: mesh_instance.mesh_asset_id().into(),
                },
                (*render_entity, *visible_entity),
                mesh_instance.current_uniform_index,
                // This example supports batching and multi draw indirect,
                // but if your pipeline doesn't support it you can use
                // `BinnedRenderPhaseType::UnbatchableMesh`
                BinnedRenderPhaseType::mesh(
                    mesh_instance.should_batch(),
                    &gpu_preprocessing_support,
                ),
            );
        }
    }
}