storyteller-bevy/studio/src/scene.rs

use bevy::{
    asset::{DependencyLoadState, LoadState, RecursiveDependencyLoadState},
    core_pipeline::{
        bloom::{BloomCompositeMode, BloomPrefilterSettings, BloomSettings},
        fxaa::Fxaa,
        tonemapping::Tonemapping,
    },
    gltf::Gltf,
    math::vec2,
    pbr::ShadowFilteringMethod,
    prelude::*,
    render::primitives::Aabb,
    scene::InstanceId,
    utils::HashSet,
};

use space_editor::{
    prelude::PrefabMarker,
    space_prefab::component::{GltfPrefab, SceneAutoChild},
};

#[cfg(feature = "wasm")]
use wasm_bindgen::prelude::wasm_bindgen;

use crate::{
    camera_controller::{CameraController, CameraControllerBundle},
    debug::DebugPoint,
    math::{Combinable, Extrema},
    MainCamera,
};

pub struct ScenePlugin;

#[derive(SystemSet, Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct SceneSpawnerSystem;

impl Plugin for ScenePlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(ClearColor(Color::hsl(0.0, 0.0, 0.25)))
            .insert_resource(Msaa::Off)
            .add_state::<SceneState>()
            .add_event::<LoadSceneEvent>()
            .add_systems(Startup, initialize)
            .add_systems(Update, init_scene_loading.in_set(SceneSpawnerSystem))
            .add_systems(
                OnEnter(SceneState::Active),
                (tag_scene_elements, frame_scene_in_view).in_set(SceneSpawnerSystem),
            );

        app.add_systems(Update, (tag_scene_elements, play_animations));

        //Auto loading screen
        //Register handles on which we will look to show loading screen
        app.init_resource::<LoadingQueue>();
        app.add_systems(
            PostUpdate,
            (
                check_loading,
                auto_add_handle::<Scene>,
                auto_add_handle::<Mesh>,
                auto_add_handle::<StandardMaterial>,
            ),
        );

        //setup default light
        app.add_systems(Startup, minimal_scene_setup);

        #[cfg(feature = "wasm")]
        app.add_systems(
            Update,
            (
                wasm::dispatch_scene_request,
                wasm::notify_state_changes
                    .map(bevy::utils::error)
                    .run_if(resource_changed::<State<SceneState>>()),
            ),
        );
    }
}

fn minimal_scene_setup(mut cmd: Commands) {
    cmd.spawn(DirectionalLightBundle {
        directional_light: DirectionalLight {
            shadows_enabled: false,
            illuminance: 20000.0,
            ..default()
        },
        transform: Transform::default().looking_at(Vec3::NEG_ONE, Vec3::Y),
        ..default()
    })
    .insert(PrefabMarker);
}

#[cfg_attr(feature = "wasm", wasm_bindgen)]
#[derive(States, Clone, Copy, Default, Debug, PartialEq, Eq, Hash)]
pub enum SceneState {
    #[default]
    None,
    Loading,
    Active,
}

#[derive(Resource, Deref)]
pub struct ActiveScene {
    #[deref]
    pub instance: InstanceId,
    pub source: Handle<Gltf>,
}

#[derive(Resource, Default)]
pub struct LoadingQueue {
    pub queue: Vec<UntypedHandle>,
}

fn check_loading(
    mut loading_queue: ResMut<LoadingQueue>,
    mut st_next: ResMut<NextState<SceneState>>,
    asset_server: Res<AssetServer>,
) {
    if !loading_queue.queue.is_empty() {
        //filter queue to only include notloaded assets

        let mut not_loaded_count = 0;

        loading_queue.queue.retain(|h| {
            if let Some((load_state, dep_load_state, rec_load_state)) =
                asset_server.get_load_states(h.id())
            {
                let is_loaded = dep_load_state == DependencyLoadState::Loaded;

                let is_failed = dep_load_state == DependencyLoadState::Failed
                    || rec_load_state == RecursiveDependencyLoadState::Failed
                    || load_state == LoadState::Failed;

                if load_state == LoadState::NotLoaded {
                    not_loaded_count += 1;
                }

                !(is_failed || is_loaded)
            } else {
                false
            }
        });

        info!("Loading {} assets", loading_queue.queue.len());

        if loading_queue.queue.len() - not_loaded_count == 0 {
            //not change state
        } else {
            st_next.set(SceneState::Loading);
        }
    } else {
        st_next.set(SceneState::Active);
    }
}

fn auto_add_handle<T: Asset>(
    mut queue: ResMut<LoadingQueue>,
    query: Query<&Handle<T>, Added<Handle<T>>>,
) {
    for h in query.iter() {
        queue.queue.push(h.clone().untyped());
    }
}

#[derive(Event, Deref)]
pub struct LoadSceneEvent(pub String);

// TODOL This could be converted into an enum that indicates what "kind" of
//       element this is (e.g. StaticMesh, SkinnedMesh, Bone, Light, Camera).
//       We don't have an immediate need for that just yet, but likely will
//       soon.
/// A tag that identifies its entity as a user-facing scene element like a prop,
/// a light, or a camera.
#[derive(Component)]
pub struct SceneElement;

#[derive(Component)]
pub struct SceneData {
    pub scene_path: String,
}

#[derive(Component)]
pub struct SceneRoot;

fn initialize(mut cmd: Commands) {
    // Camera, graphics settings
    let focus = Vec3::ZERO;
    let camera_xform = Transform::from_xyz(-2.5, 4.5, 5.0).looking_at(focus, Vec3::Y);

    cmd.spawn((
        Name::new("Main Camera"),
        MainCamera,
        Camera3dBundle {
            transform: camera_xform,
            tonemapping: Tonemapping::BlenderFilmic,
            camera: Camera {
                hdr: true,
                ..default()
            },
            ..default()
        },
        CameraControllerBundle::new(focus, camera_xform.translation),
        ShadowFilteringMethod::Hardware2x2,
        BloomSettings {
            composite_mode: BloomCompositeMode::EnergyConserving,
            intensity: 0.25,
            prefilter_settings: BloomPrefilterSettings {
                threshold: 0.8,
                threshold_softness: 0.5,
            },
            ..default()
        },
        Fxaa::default(),
    ));
}

fn init_scene_loading(mut cmd: Commands, mut er: EventReader<LoadSceneEvent>) {
    if let Some(id) = er.read().last() {
        info!("Loading scene: {}", id.0);
        cmd.spawn((
            SpatialBundle::default(),
            space_editor::prelude::GltfPrefab {
                path: format!("gltf/{}", id.0),
                scene: "Scene0".to_string(),
            },
            PrefabMarker,
            SceneRoot,
        ));
    }
}

#[derive(Component)]
pub struct NoMeshElement;

#[derive(Component)]
pub struct InitedKointSkinnedMesh;

fn tag_scene_elements(
    mut cmd: Commands,
    // A "glTF mesh" is actually a container for one or more "glTF primitives",
    // which is basically a bundle of vertices paired with the material used to
    // shade them. It's those "sub-meshes" that carry the actual mesh handles,
    // but we don't necessarily want to expose those to users a la carte,
    // because they often consist of stuff like a single eyeball, or just the
    // leafy parts of a plant. Here we only want to tag the elements that a user
    // might want to individually address when arranging or animating a scene.
    q_named_non_primitives: Query<
        (Entity, &Children),
        (
            With<SceneAutoChild>,
            With<Name>,
            Without<Handle<Mesh>>,
            Without<SceneElement>,
            Without<NoMeshElement>,
        ),
    >,
    mesh_entities: Query<Entity, With<Handle<Mesh>>>,
) {
    for (e, children) in q_named_non_primitives.iter() {
        let mut has_mesh = false;
        for child in children {
            if mesh_entities.contains(*child) {
                has_mesh = true;
                break;
            }
        }
        if has_mesh {
            cmd.entity(e).insert(SceneElement);
        } else {
            cmd.entity(e).insert(NoMeshElement);
        }
    }
}

fn frame_scene_in_view(
    mut cmd: Commands,
    q_aabbs: Query<(&Aabb, &GlobalTransform), With<SceneAutoChild>>,
    mut q_camera: Query<(&mut CameraController, &mut Transform, &Projection)>,
) {
    let Ok((mut cam, mut cam_xform, cam_proj)) = q_camera.get_single_mut() else {
        error!("Failed to get camera controller");
        return;
    };

    let Some(aabb) = Aabb::combine(q_aabbs.iter()) else {
        return;
    };

    let world2cam = cam_xform.compute_matrix().inverse();
    let plane_origin = Vec3::from(aabb.center);
    let plane_normal = cam_xform.forward();

    let Some((_depth, bounding_rect)) = aabb
        .extrema()
        .iter()
        .copied()
        .map(|mut p| {
            cmd.spawn(DebugPoint {
                point: p,
                color: Color::CYAN,
            });

            // Inflate the combined AABB to give the model some breathing room
            p *= 1.5;

            // Project each corner onto the camera plane
            let v = p - plane_origin;
            let proj = p - v.project_onto_normalized(plane_normal);

            cmd.spawn(DebugPoint {
                point: proj,
                color: Color::FUCHSIA,
            });

            // Transform into the camera's local space
            world2cam.transform_point(proj)
        })
        // Fold the camera-space plane-projected points into a bounding rect and view depth
        .fold(None, |accum: Option<(f32, Rect)>, p: Vec3| match accum {
            Some((depth, rect)) => Some((depth, rect.union_point(vec2(p.x, p.y)))),
            None => {
                let depth = p.z.abs();
                let p = vec2(p.x, p.y);
                Some((depth, Rect::from_corners(p, p)))
            }
        })
    else {
        return;
    };

    let pproj = match cam_proj {
        Projection::Perspective(p) => p,
        _ => panic!("Expected a perspective projection"),
    };

    // Formula for the target camera distance
    // (c/o https://stackoverflow.com/questions/2866350/move-camera-to-fit-3d-scene):
    //
    // d = (s/2) / tan(a/2)
    // where:
    //   d = distance
    //   s = size of the largest camera-aligned axis in world coords
    //   a = field-of-view angle for the same axis
    let Vec2 {
        x: width,
        y: height,
    } = bounding_rect.size();

    let s = f32::max(width, height);
    let a = if s == height {
        pproj.fov
    } else {
        f32::atan(pproj.aspect_ratio * f32::tan(pproj.fov / 2.))
    };

    // TODO:
    // According to "I eat babies" at the Ars Technica forums
    // (https://arstechnica.com/civis/threads/if-you-have-aspect-ratio-and-vertical-fov-can-you-calc-horizontal-fov.37447/):
    //
    // tan(fov_v / 2) = y / z
    // tan(fov_h / 2) = x / z
    //
    // Theoretically, this would let us skip some costly tan/atan
    // computations, but I must be missing something, because just
    // substituting (s / depth) for the divisor below would make the
    // whole equation independent of the FOV, which is obviously
    // wrong.
    //
    // Unfortunately I'm embarrassingly bad at algebra, but there
    // must be some room for simplification here.
    let radius = (s / 2.) / f32::tan(a / 2.);

    cam.focus = aabb.center.into();
    cam.radius = radius;
    cam_xform.translation = cam.focus + cam_xform.back() * radius;
}

fn play_animations(
    mut cmd: Commands,
    mut l_done: Local<HashSet<String>>,
    r_assets: Res<AssetServer>,
    ra_gltfs: Res<Assets<Gltf>>,
    ra_anims: Res<Assets<AnimationClip>>,
    q_prefabs: Query<&GltfPrefab>,
    q_named_ents: Query<(Entity, &Name)>,
    mut q_players: Query<&mut AnimationPlayer>,
) {
    for prefab in q_prefabs.iter() {
        if l_done.contains(&prefab.path) {
            continue;
        }

        // TODO: Get the `Handle<Gltf>` directly from the entity once it's available
        let gltf_handle = match r_assets.get_handle::<Gltf>(&prefab.path) {
            Some(handle) => {
                info!("Returning existing handle");
                handle
            }
            None => {
                info!("Loading from path");
                r_assets.load::<Gltf>(&prefab.path)
            }
        };

        let Some(gltf) = ra_gltfs.get(gltf_handle) else {
            continue;
        };

        let mut playing_anims = 0_usize;
        for anim_handle in gltf.animations.iter() {
            let Some(anim) = ra_anims.get(anim_handle.clone()) else {
                continue;
            };

            let Some(target) = q_named_ents.iter().find_map(|(ent, name)| {
                if anim.compatible_with(name) {
                    Some(ent)
                } else {
                    None
                }
            }) else {
                continue;
            };

            if let Ok(mut player) = q_players.get_mut(target) {
                if !player.is_playing_clip(anim_handle) {
                    player.play(anim_handle.clone()).repeat();
                }
            } else {
                let mut player = AnimationPlayer::default();
                player.play(anim_handle.clone()).repeat();
                cmd.entity(target).insert(player);
            }

            playing_anims += 1;
        }

        if playing_anims == gltf.animations.len() {
            l_done.insert(prefab.path.clone());
        }
    }
}

#[cfg(feature = "wasm")]
mod wasm {
    use bevy::{prelude::*, window::PrimaryWindow};
    use wasm_bindgen::prelude::*;

    use super::{LoadSceneEvent, SceneState};
    use crate::wasm::{self, StaticBuffer, WasmMessageBuffer};

    static OBJECT_LOAD_REQUEST: WasmMessageBuffer<String> = WasmMessageBuffer::new();

    #[wasm_bindgen(js_name = loadObject)]
    pub fn load_object(id: &str) {
        OBJECT_LOAD_REQUEST.write_message(id.into());
    }

    // FIXME: Is it an object or a scene? Inconsistent terminology is
    //        unnecessarily confusing
    pub(super) fn dispatch_scene_request(mut ew_load_scene: EventWriter<LoadSceneEvent>) {
        if let Some(id) = OBJECT_LOAD_REQUEST.take_message() {
            ew_load_scene.send(LoadSceneEvent(id));
        }
    }

    #[wasm_bindgen(typescript_custom_section)]
    const SCENE_STATE_EVENT: &str = r#"
export interface SceneStateEvent extends CustomEvent {
    type: "scene-state";
    detail: SceneState;
}

declare global {
    export interface HTMLElementEventMap {
        "scene-state": SceneStateEvent;
    }
}
    "#;

    pub(super) fn notify_state_changes(
        r_state: Res<State<SceneState>>,
        q_window: Query<&Window, With<PrimaryWindow>>,
    ) -> Result<(), String> {
        if let Ok(window) = q_window.get_single() {
            wasm::dispatch_to_js(window, "scene-state", &(**r_state as u32))?;
        }
        Ok(())
    }
}