monologue-trees/bin/gltf-inspect.rs

use bevy::{
    core_pipeline::clear_color::ClearColorConfig,
    gltf::Gltf,
    input::{
        keyboard::KeyboardInput,
        mouse::{MouseMotion, MouseWheel},
        ButtonState,
    },
    pbr::CascadeShadowConfigBuilder,
    prelude::*,
    render::{
        camera::RenderTarget,
        render_resource::{
            Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
        },
    },
    utils::HashSet,
};

fn main() {
    App::new()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "GLTF Inspector".into(),
                resolution: (640., 480.).into(),
                ..default()
            }),
            ..default()
        }))
        .add_event::<ManageActive>()
        .add_startup_system(load_models)
        .add_startup_system(spawn_base_scene)
        .add_startup_system(spawn_base_ui)
        .add_system(spawn_models)
        .add_system(spawn_ui)
        .add_system(catalog_animations)
        .add_system(play_animation)
        .add_system(rotate_model)
        .add_system(zoom_model)
        .add_system(scroll)
        .add_system(select)
        .add_system(manage_active)
        .run();
}

///
/// Stores GLTF handles for later use
#[derive(Resource)]
struct Models {
    handles: Vec<Handle<Gltf>>,
}

///
/// Marks GLTF model as inspectable
#[derive(Component)]
struct Inspect;

#[derive(Component)]
struct SelectionUI;

#[derive(Component)]
struct PreviewCamera;

#[derive(Component)]
struct ScrollingList;

#[derive(Component)]
struct Preview(Handle<Image>);

#[derive(Component)]
struct Container;

#[derive(Component)]
struct Active;

///
/// Event for managing which entities are tagged "Active"
#[derive(Debug)]
struct ManageActive(Option<Entity>);

///
/// Load all GLTF models on startup
fn load_models(mut commands: Commands, ass: Res<AssetServer>) {
    let weak_handles = ass.load_folder("models").expect("Load gltfs");
    let handles: Vec<Handle<Gltf>> = weak_handles
        .iter()
        .map(|weak| weak.clone().typed::<Gltf>())
        .collect();
    // info!("Scene Handles: {:#?}", handles);
    commands.insert_resource(Models { handles });
}

///
/// Spawn base scene
fn spawn_base_scene(mut commands: Commands) {
    commands.spawn((
        Camera2dBundle {
            camera_2d: Camera2d {
                clear_color: ClearColorConfig::Custom(Color::BLACK),
            },
            ..default()
        },
        UiCameraConfig { ..default() },
        SelectionUI,
    ));
}

fn spawn_base_ui(mut commands: Commands) {
    commands
        .spawn((
            NodeBundle {
                style: Style {
                    justify_content: JustifyContent::Center,
                    size: Size::all(Val::Percent(90.0)),
                    overflow: Overflow::Hidden,
                    ..default()
                },
                ..default()
            },
            SelectionUI,
        ))
        .with_children(|parent| {
            parent.spawn((
                NodeBundle {
                    style: Style {
                        flex_wrap: FlexWrap::Wrap,
                        flex_direction: FlexDirection::Row,
                        justify_content: JustifyContent::SpaceAround,
                        size: Size::AUTO,
                        max_size: Size::UNDEFINED,
                        ..default()
                    },
                    ..default()
                },
                ScrollingList,
            ));
        });
}

///
/// Spawn a loaded scene for inspection
/// TODO: Update/add/delete when files are updated
fn spawn_models(
    mut commands: Commands,
    mut scene_evr: EventReader<AssetEvent<Scene>>,
    mut images: ResMut<Assets<Image>>,
) {
    if !scene_evr.is_empty() {
        info!("Spawning scenes");
    }

    for ev in scene_evr.iter() {
        match ev {
            AssetEvent::Created { handle } => {
                info!("Creating scene {:#?}", handle);

                // Preview image
                let preview_image_handle = {
                    info!("Creating preview");

                    let size = Extent3d {
                        width: 256,
                        height: 256,
                        ..default()
                    };

                    // Create render target image for the preview camera
                    let mut image = Image {
                        texture_descriptor: TextureDescriptor {
                            label: None,
                            size,
                            dimension: TextureDimension::D2,
                            format: TextureFormat::Bgra8UnormSrgb,
                            mip_level_count: 1,
                            sample_count: 1,
                            usage: TextureUsages::TEXTURE_BINDING
                                | TextureUsages::COPY_DST
                                | TextureUsages::RENDER_ATTACHMENT,
                            view_formats: &[],
                        },
                        ..default()
                    };

                    // Fill with zeroes
                    image.resize(size);

                    let image_handle = images.add(image);

                    image_handle
                };

                let local = {
                    // Get a unique number for this resource from the handle
                    let idx = handle.id().reflect_hash().unwrap();
                    // Set the origin of this scene to [idx, idx, idx];
                    let origin = Vec3::ONE * ((idx % 1000) as f32);
                    // Transform pointing at origin
                    Transform::from_translation(origin)
                };

                // Spawn the actual scene
                commands
                    .spawn((
                        SpatialBundle {
                            transform: local,
                            ..default()
                        },
                        Inspect,
                        Container,
                        Preview(preview_image_handle.clone()),
                    ))
                    .with_children(|builder| {
                        let camera_location =
                            Transform::from_xyz(0.0, 0.0, 10.0).looking_at(Vec3::ZERO, Vec3::Y);

                        // Spawn preview camera
                        builder.spawn((
                            Camera3dBundle {
                                camera_3d: Camera3d {
                                    clear_color: ClearColorConfig::Custom(Color::WHITE),
                                    ..default()
                                },
                                camera: Camera {
                                    order: -1,
                                    target: RenderTarget::Image(preview_image_handle.clone()),
                                    ..default()
                                },
                                transform: camera_location,
                                ..default()
                            },
                            UiCameraConfig { show_ui: false },
                            PreviewCamera,
                        ));

                        // Spawn window camera
                        builder.spawn((
                            Camera3dBundle {
                                camera_3d: Camera3d {
                                    clear_color: ClearColorConfig::Custom(Color::WHITE),
                                    ..default()
                                },
                                camera: Camera {
                                    is_active: false,
                                    ..default()
                                },
                                transform: camera_location,
                                ..default()
                            },
                            Inspect,
                            Preview(preview_image_handle.clone()),
                        ));

                        builder.spawn((
                            DirectionalLightBundle {
                                directional_light: DirectionalLight {
                                    shadows_enabled: true,
                                    ..default()
                                },
                                cascade_shadow_config: CascadeShadowConfigBuilder {
                                    num_cascades: 1,
                                    maximum_distance: 1.6,
                                    ..default()
                                }
                                .into(),
                                ..default()
                            },
                            Inspect,
                        ));

                        builder.spawn((
                            SceneBundle {
                                scene: handle.clone(),
                                ..default()
                            },
                            Inspect,
                        ));
                    });
            }
            AssetEvent::Removed { .. } => {
                todo!("Remove deleted scene")
            }
            AssetEvent::Modified { .. } => {
                todo!("Update modified scene")
            }
        }
    }
}

fn spawn_ui(
    mut commands: Commands,
    query: Query<Entity, With<ScrollingList>>,
    previews: Query<&Preview, (Added<Preview>, Without<Camera>, Without<Interaction>)>,
) {
    // UI container
    if let Ok(scrolling_list_container) = query.get_single() {
        let mut entity_commands = commands.entity(scrolling_list_container);
        entity_commands.with_children(|parent| {
            for Preview(image_handle) in previews.iter() {
                // Preview Image
                info!("Spawning image preview");
                parent
                    .spawn((
                        ButtonBundle { ..default() },
                        SelectionUI,
                        Preview(image_handle.clone()),
                    ))
                    .with_children(|parent| {
                        parent.spawn(ImageBundle {
                            style: Style {
                                size: Size::all(Val::Px(256.0)),
                                padding: UiRect::all(Val::Px(5.0)),
                                ..default()
                            },
                            image: UiImage {
                                texture: image_handle.clone(),
                                ..default()
                            },
                            ..default()
                        });
                    });
            }
        });
    }
}

fn catalog_animations(
    mut events: EventReader<AssetEvent<AnimationClip>>,
    clips: Res<Assets<AnimationClip>>,
) {
    for event in events.iter() {
        match event {
            AssetEvent::Created { handle } => {
                if let Some(clip) = clips.get(handle) {
                    info!("Event: {:#?}", clip.curves());
                }
            }
            _ => (),
        }
    }
}

fn play_animation(
    mut events: EventReader<ManageActive>,
    mut players: Query<&mut AnimationPlayer>,
    children: Query<&Children>,
    models: Res<Models>,
    gltfs: Res<Assets<Gltf>>,
) {
    for event in events.iter() {
        match event {
            ManageActive(None) => {
                info!("Disabling all animations");
                for mut player in players.iter_mut() {
                    player.stop_repeating();
                }
            }
            ManageActive(Some(entity)) => {
                info!("Setting {:?} Active Animation", entity);
                for child in children.iter_descendants(*entity) {
                    if let Ok(mut player) = players.get_mut(child) {
                        for gltf_handle in models.handles.iter() {
                            if let Some(gltf) = gltfs.get(gltf_handle) {
                                for animation_handle in gltf.animations.iter() {
                                    player.start(animation_handle.clone()).repeat();
                                }
                            } else {
                                info!("Failed to get GLTF handle");
                            }
                        }
                    } else {
                        info!("Failed to get active animation player");
                    }
                }
            }
        }
    }
}

///
/// Rotate a model as part of inspection
/// TODO: move light with model (want to see shadows)
fn rotate_model(
    buttons: Res<Input<MouseButton>>,
    mut mouse_evr: EventReader<MouseMotion>,
    mut transforms: Query<&mut Transform, (With<Inspect>, With<Active>, Without<Camera>)>,
) {
    if buttons.pressed(MouseButton::Left) {
        for MouseMotion { delta } in mouse_evr.iter() {
            for mut transform in transforms.iter_mut() {
                let rot_y = delta.x / 1000.0;
                let rot_x = delta.y / 1000.0;
                transform.rotate_y(rot_y);
                transform.rotate_x(rot_x);
            }
        }
    }
}

///
/// Zoom in and out of the model
/// TODO: Only modify selected entities
fn zoom_model(
    keys: Res<Input<KeyCode>>,
    mut wheel_evr: EventReader<MouseWheel>,
    mut transforms: Query<&mut Transform, (With<Inspect>, With<Active>, Without<Camera>)>,
) {
    if keys.pressed(KeyCode::LShift) {
        for ev in wheel_evr.iter() {
            for mut transform in transforms.iter_mut() {
                let scale = (Vec3::ONE * ev.y) / 100.0;
                transform.scale += scale;
            }
        }
    }
}

fn scroll(
    mut scroll_evr: EventReader<MouseWheel>,
    mut query: Query<&mut Style, With<ScrollingList>>,
) {
    for ev in scroll_evr.iter() {
        for mut s in query.iter_mut() {
            s.position.top = match s.position.top {
                Val::Px(current) => Val::Px(current + (ev.y * 5.0)),
                _ => Val::Px(0.0),
            };
        }
    }
}

///
/// Click a UI element to select
///
/// This is a really ugly implementation. I'm not really happy with how we have to navigate the ECS
/// parent/child hierarchy. There should be a more direct way to correlate the scene with the
/// button.
fn select(
    query: Query<(&Interaction, &Preview), (With<SelectionUI>, Changed<Interaction>)>,
    mut selection_ui: Query<&mut Visibility, (With<SelectionUI>, Without<Parent>)>,
    mut ui_camera: Query<&mut Camera, (With<SelectionUI>, Without<Inspect>)>,
    mut scene_camera: Query<(Entity, &mut Camera, &Preview), (With<Inspect>, Without<SelectionUI>)>,
    mut key_evr: EventReader<KeyboardInput>,
    mut selected: Local<Option<Entity>>, // Active camera index
    parent_search: Query<&Children>,
    parents: Query<Entity, (With<Children>, Without<Parent>)>, // TODO: Constrain
    mut events: EventWriter<ManageActive>,
) {
    for (interaction, selected_preview) in query.iter() {
        if interaction == &Interaction::Clicked {
            // Hide UI
            let mut ui_vis = selection_ui.single_mut();
            *ui_vis = Visibility::Hidden;

            // Disable UI camera
            let mut ui_cam = ui_camera.single_mut();
            ui_cam.is_active = false;

            // Determine selected scene
            *selected = scene_camera
                .iter()
                .find(|(_, _, preview)| selected_preview.0 == preview.0)
                .map(|(entity, _, _)| entity);

            // Enable scene camera
            let (_, mut scene_cam, _) = scene_camera
                .get_mut(selected.expect("Selected scene should be set"))
                .expect("Failed to get Scene camera");
            scene_cam.is_active = true;

            // Set relevant entities active
            let message = parents.iter().find(|&parent| {
                parent_search
                    .iter_descendants(parent)
                    .find(|&entity| Some(entity) == *selected)
                    .is_some()
            });
            events.send(ManageActive(message));
        }
    }
    for ev in key_evr.iter() {
        match ev {
            KeyboardInput {
                state: ButtonState::Pressed,
                key_code: Some(KeyCode::Escape),
                ..
            } => {
                if let Some(s) = *selected {
                    // Set all inactive
                    events.send(ManageActive(None));

                    // Disable scene camera
                    let (_, mut scene_cam, _) =
                        scene_camera.get_mut(s).expect("Failed to get Scene camera");
                    scene_cam.is_active = false;

                    // Enable UI camera
                    let mut ui_cam = ui_camera.single_mut();
                    ui_cam.is_active = true;

                    // Make UI visible
                    let mut ui_vis = selection_ui.single_mut();
                    *ui_vis = Visibility::Inherited;
                }
            }
            _ => (),
        }
    }
}

fn manage_active(
    mut commands: Commands,
    mut events: EventReader<ManageActive>,
    query: Query<&Children>,
    current: Query<Entity, With<Active>>,
) {
    for event in events.iter() {
        info!("Setting active: {:?}", event);
        match event {
            ManageActive(None) => {
                for entity in current.iter() {
                    if let Some(mut entity_commands) = commands.get_entity(entity) {
                        entity_commands.remove::<Active>();
                    }
                }
            }
            ManageActive(Some(entity)) => {
                for child in query.iter_descendants(*entity) {
                    if let Some(mut child_commands) = commands.get_entity(child) {
                        // info!(
                        //     "Active entity components: {:?} {:?}",
                        //     child,
                        //     names.get(child)
                        // );
                        // child_commands.log_components();
                        child_commands.insert(Active);
                    }
                }
            }
        }
    }
}