snow_trail/src/bundles/terrain.rs

use std::rc::Rc;

use glam::Vec2;
use nalgebra::vector;
use rapier3d::prelude::{ColliderBuilder, RigidBodyBuilder};

use crate::components::{MeshComponent, PhysicsComponent, TreeInstancesComponent};
use crate::entity::EntityHandle;
use crate::loaders::mesh::{InstanceData, InstanceRaw, Mesh};
use crate::loaders::terrain::load_heightfield_from_exr;
use crate::physics::PhysicsManager;
use crate::render;
use crate::world::{Transform, World};

pub struct TerrainConfig
{
    pub gltf_path: String,
    pub heightmap_path: String,
    pub size: Vec2,
}

impl TerrainConfig
{
    pub fn new(gltf_path: &str, heightmap_path: &str, size: Vec2) -> Self
    {
        Self {
            gltf_path: gltf_path.to_string(),
            heightmap_path: heightmap_path.to_string(),
            size,
        }
    }

    pub fn default() -> Self
    {
        Self {
            gltf_path: "meshes/terrain.gltf".to_string(),
            heightmap_path: "textures/terrain_heightmap.exr".to_string(),
            size: Vec2::new(1000.0, 1000.0),
        }
    }
}

pub struct TerrainBundle;

impl TerrainBundle
{
    pub fn spawn(
        world: &mut World,
        mesh_data: Vec<(Mesh, Vec<InstanceData>)>,
        config: &TerrainConfig,
    ) -> anyhow::Result<EntityHandle>
    {
        let transform = Transform::IDENTITY;
        let mut first_entity = None;
        let mut physics_added = false;

        for (mesh, instances) in mesh_data
        {
            let entity = world.spawn();

            if first_entity.is_none()
            {
                first_entity = Some(entity);
            }

            if instances.is_empty()
            {
                world.transforms.insert(entity, transform);
                world.meshes.insert(
                    entity,
                    MeshComponent {
                        mesh: Rc::new(mesh),
                        pipeline: render::Pipeline::Standard,
                        instance_buffer: None,
                        num_instances: 1,
                        tile_scale: 2.0,
                        enable_dissolve: false,
                        enable_snow_light: true,
                    },
                );
                world.names.insert(entity, "Terrain".to_string());

                if !physics_added
                {
                    let heights = load_heightfield_from_exr(&config.heightmap_path)?;

                    let height_scale = 1.0;
                    let scale = vector![config.size.x, height_scale, config.size.y];

                    let body = RigidBodyBuilder::fixed()
                        .translation(transform.get_position().into())
                        .build();

                    let rigidbody_handle = PhysicsManager::add_rigidbody(body);

                    let collider = ColliderBuilder::heightfield(heights.clone(), scale).build();

                    let collider_handle =
                        PhysicsManager::add_collider(collider, Some(rigidbody_handle));

                    PhysicsManager::set_heightfield_data(
                        heights,
                        scale,
                        transform.get_position().into(),
                    );

                    world.physics.insert(
                        entity,
                        PhysicsComponent {
                            rigidbody: rigidbody_handle,
                            collider: Some(collider_handle),
                        },
                    );

                    physics_added = true;
                }
            }
            else
            {
                let num_instances = instances.len();

                let instance_raw: Vec<InstanceRaw> = instances.iter().map(|i| i.to_raw()).collect();

                let instance_buffer = render::with_device(|device| {
                    use wgpu::util::DeviceExt;
                    device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("Tree Instance Buffer"),
                        contents: bytemuck::cast_slice(&instance_raw),
                        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                    })
                });

                world.transforms.insert(entity, Transform::IDENTITY);
                world.meshes.insert(
                    entity,
                    MeshComponent {
                        mesh: Rc::new(mesh),
                        pipeline: render::Pipeline::Standard,
                        instance_buffer: Some(instance_buffer.clone()),
                        num_instances: num_instances as u32,
                        tile_scale: 4.0,
                        enable_dissolve: true,
                        enable_snow_light: false,
                    },
                );
                world.tree_instances.insert(
                    entity,
                    TreeInstancesComponent::new(instances, Rc::new(instance_buffer)),
                );
                world.names.insert(entity, "Trees".to_string());
            }
        }

        first_entity.ok_or_else(|| anyhow::anyhow!("No meshes found in glTF file"))
    }
}