martian-chess/src/display3d.rs

// TODO: Unify the selection logic
// If you select in 2d, it should "be" selected in 3d.
// And if you select in 3d, 2d is selected.
use crate::{
    game::{Board, BoardIndex, Piece, Side},
    prelude::*,
};
use bevy::{
    core_pipeline::Skybox,
    ecs::removal_detection::RemovedComponentReader,
    input::mouse::{MouseButtonInput, MouseMotion, MouseWheel},
    render::render_resource::{TextureViewDescriptor, TextureViewDimension},
    window::PrimaryWindow,
};

pub(crate) struct Display3dPlugin;

impl Plugin for Display3dPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(OnEnter(GameState::Loading), load_assets)
            .add_systems(
                OnExit(GameState::Loading),
                (initialize, fix_skybox.before(initialize)),
            )
            .add_systems(
                Update,
                (
                    menu::exit_to_menu.run_if(in_state(GameState::Display3d)),
                    set_piece_model.run_if(any_component_added::<Piece>),
                    set_board_model.run_if(any_component_added::<game::BoardComponent>),
                    set_piece_position.run_if(any_component_changed::<BoardIndex>),
                    set_piece_texture.run_if(any_component_changed::<Side>),
                    select_3d
                        .run_if(in_state(GameState::Display3d))
                        .run_if(on_event::<MouseButtonInput>()),
                    pick_up.run_if(any_component_added::<game::Selected>),
                    put_down.run_if(any_component_removed::<game::Selected>()),
                ),
            )
            .add_systems(
                Update,
                (
                    move_camera
                        .run_if(in_state(GameState::Display3d))
                        .run_if(on_event::<MouseMotion>()),
                    gizmo_system.run_if(in_state(GameState::Display3d)),
                    mouse_zoom
                        .run_if(in_state(GameState::Display3d))
                        .run_if(on_event::<MouseWheel>()),
                    selected_gizmo,
                )
                    .run_if(resource_exists::<debug::DebugEnabled>()),
            )
            .add_systems(
                OnEnter(GameState::Display3d),
                (activate::<Display3d>, set_piece_texture),
            )
            .add_systems(OnExit(GameState::Display3d), deactivate::<Display3d>);
    }
}

#[derive(Debug, Component)]
pub(crate) struct Display3d;

#[derive(Debug, Resource)]
struct AssetsMap {
    models: Handle<Gltf>,
    skybox: Handle<Image>,
}

/// Load 3d models
/// This is kind of pulling double duty.
/// Both loads the GLTF file _and_ populates the ModelMap once that is loaded.
fn load_assets(server: Res<AssetServer>, mut commands: Commands) {
    commands.insert_resource(AssetsMap {
        models: server.load("models/Martian Chess.glb"),
        skybox: server.load("images/skybox.png"),
    });
}

/// Initialize the 3d board
fn initialize(mut commands: Commands, board: Res<game::Board>, assets: Res<AssetsMap>) {
    info!("Initialize 3d camera");
    // let handle = server.load("images/mars.hdr");
    commands.spawn((
        Display3d,
        Camera3dBundle {
            camera: Camera {
                is_active: false,
                hdr: true,
                ..default()
            },
            transform: Transform::from_xyz(0.0, 20.0, 10.0).looking_at(Vec3::ZERO, Vec3::Y),
            ..default()
        },
        Skybox(assets.skybox.clone()),
        EnvironmentMapLight {
            diffuse_map: assets.skybox.clone(),
            specular_map: assets.skybox.clone(),
        },
        UiCameraConfig { show_ui: true },
    ));

    info!("Initializing root");
    commands
        .spawn((
            SpatialBundle {
                visibility: Visibility::Hidden,
                ..default()
            },
            Display3d,
        ))
        .with_children(|parent| {
            info!("Initializing 3D lights!");
            parent.spawn((
                Display3d,
                PointLightBundle {
                    point_light: PointLight {
                        intensity: 3000.0,
                        shadows_enabled: true,
                        ..default()
                    },
                    transform: Transform::from_xyz(0.0, 10.0, 5.0),
                    ..default()
                },
            ));

            info!("Intializing 3D Board!");
            parent
                .spawn((Display3d, game::BoardComponent, SceneBundle { ..default() }))
                .with_children(|parent| {
                    board.pieces().iter().for_each(|(index, piece)| {
                        let side = Board::side(index).expect("Spawn valid side");

                        parent.spawn((
                            side,
                            Display3d,
                            piece.clone(),
                            index.clone(),
                            SceneBundle { ..default() },
                        ));
                    });
                });
        });
}

fn fix_skybox(mut images: ResMut<Assets<Image>>, assets: Res<AssetsMap>) {
    let image = images.get_mut(&assets.skybox).unwrap();
    info!("Loaded skybox image");
    // NOTE: PNGs do not have any metadata that could indicate they contain a cubemap texture,
    // so they appear as one texture. The following code reconfigures the texture as necessary.
    if image.texture_descriptor.array_layer_count() == 1 {
        image.reinterpret_stacked_2d_as_array(
            image.texture_descriptor.size.height / image.texture_descriptor.size.width,
        );
        image.texture_view_descriptor = Some(TextureViewDescriptor {
            dimension: Some(TextureViewDimension::Cube),
            ..default()
        });
    }
}

/// Set the model for each piece based on the game::Piece::* marker
fn set_piece_model(
    mut events: Query<(&mut Handle<Scene>, &Piece), (Added<game::Piece>, With<Display3d>)>,
    assets_map: Res<AssetsMap>,
    gltfs: Res<Assets<Gltf>>,
) {
    events.iter_mut().for_each(|(mut handle, piece)| {
        let gltf = gltfs.get(&assets_map.models).expect("Load GLTF content");
        *handle = match piece {
            game::Piece::Pawn => gltf.named_scenes.get("Pawn"),
            game::Piece::Drone => gltf.named_scenes.get("Drone"),
            game::Piece::Queen => gltf.named_scenes.get("Queen"),
        }
        .expect("Game board model")
        .clone();
    });
}

fn set_board_model(
    mut events: Query<&mut Handle<Scene>, (With<game::BoardComponent>, With<Display3d>)>,
    assets_map: Res<AssetsMap>,
    gltfs: Res<Assets<Gltf>>,
) {
    events.iter_mut().for_each(|mut handle| {
        let gltf = gltfs.get(&assets_map.models).expect("Load GLTF content");
        *handle = gltf
            .named_scenes
            .get("Gameboard")
            .expect("Game board model")
            .clone();
    })
}

/// Sets a piece location given it's board index
fn set_piece_position(
    mut events: Query<
        (&mut Transform, &BoardIndex),
        (With<game::Piece>, With<Display3d>, Changed<BoardIndex>),
    >,
) {
    events.iter_mut().for_each(|(mut t, i)| {
        t.translation = board_translation(i);
    });
}

/// Given a board index returns the Vec3 location in space
fn board_translation(&BoardIndex { x, y }: &BoardIndex) -> Vec3 {
    // Scale x down by 4 to account for -4..4 scaling
    let x = x as i8 - 4;
    // Mirror y axis because our board index is inverted...
    let y = -(y as i8) + 1;

    let x = if x < 0 {
        x as f32 * 1.3 + 0.325
    } else {
        x as f32 * 1.3 + 1.0
    };

    let y = y as f32 * 1.3 + 0.65;

    Vec3::new(x, 0.0, y)
}

fn gizmo_system(mut gizmos: Gizmos) {
    for y in 0..4 {
        for x in 0..8 {
            gizmos.cuboid(
                Transform::from_translation(board_translation(&BoardIndex { x, y }))
                    .with_scale(Vec3::splat(1.25)),
                Color::PURPLE,
            )
        }
    }
}

/// TODO: This has bad feel, needs to be tuned
fn move_camera(
    buttons: Res<Input<MouseButton>>,
    mut events: EventReader<MouseMotion>,
    mut camera: Query<&mut Transform, (With<Display3d>, With<Camera>)>,
) {
    events.iter().for_each(|MouseMotion { delta }| {
        if buttons.pressed(MouseButton::Left) {
            camera.iter_mut().for_each(|mut t| {
                t.rotate_around(Vec3::ZERO, Quat::from_rotation_y(delta.x / 256.0));
                t.rotate_around(Vec3::ZERO, Quat::from_rotation_x(delta.y / 256.0));
                t.look_at(Vec3::ZERO, Vec3::Y);
            });
        }
    });
}

fn mouse_zoom(
    mut events: EventReader<MouseWheel>,
    mut camera: Query<&mut Transform, (With<Display3d>, With<Camera>)>,
) {
    events.iter().for_each(|MouseWheel { y, .. }| {
        camera.iter_mut().for_each(|mut t| {
            t.translation *= 1.0 - (*y / 4.0);
            t.look_at(Vec3::ZERO, Vec3::Y);
        });
    });
}

/// Set the Texture for a piece given it's position (left or right) on the bord.
/// Executed when Side is changed or upon entry to Display3d state
/// Getting this to run _after_ GLTF is loaded is a pain.
/// PERF: We are saving what to work on in a Vector which is bad.
/// CAVEAT: We are only exeucting this when a piece changes or state is changed.
fn set_piece_texture(
    events: Query<(Entity, &Piece, &Side), (With<game::Piece>, With<Display3d>, Changed<Side>)>,
    all: Query<(Entity, &Piece, &Side), (With<game::Piece>, With<Display3d>)>,
    gltfs: Res<Assets<Gltf>>,
    assets_map: Res<AssetsMap>,
    children: Query<&Children>,
    mut models: Query<(&Name, &mut Handle<StandardMaterial>)>,
) {
    let pieces = if events.is_empty() {
        all.iter().collect::<Vec<(Entity, &Piece, &Side)>>()
    } else {
        events.iter().collect::<Vec<(Entity, &Piece, &Side)>>()
    };
    pieces.iter().for_each(|(entity, piece, side)| {
        if let Some(gltf) = gltfs.get(&assets_map.models) {
            children.iter_descendants(*entity).for_each(|child| {
                if let Ok((n, mut m)) = models.get_mut(child) {
                    match (*piece, *side, n.as_str()) {
                        (Piece::Queen, Side::A, "Queen.0") => {
                            *m = gltf
                                .named_materials
                                .get("Queen")
                                .expect("Load Red Queen texture")
                                .clone()
                        }
                        (Piece::Queen, Side::A, "Queen.1") => {
                            *m = gltf
                                .named_materials
                                .get("Dots")
                                .expect("Load Red Dots texture")
                                .clone()
                        }
                        (Piece::Queen, Side::B, "Queen.0") => {
                            *m = gltf
                                .named_materials
                                .get("QueenBlue")
                                .expect("Load Blue Queen texture")
                                .clone()
                        }
                        (Piece::Queen, Side::B, "Queen.1") => {
                            *m = gltf
                                .named_materials
                                .get("DotsBlue")
                                .expect("Load Red Dots texture")
                                .clone()
                        }
                        (Piece::Drone, Side::A, "Drone.0") => {
                            *m = gltf
                                .named_materials
                                .get("Drone")
                                .expect("Load Red Drone texture")
                                .clone()
                        }
                        (Piece::Drone, Side::A, "Drone.1") => {
                            *m = gltf
                                .named_materials
                                .get("Dots")
                                .expect("Load Red Dots texture")
                                .clone()
                        }
                        (Piece::Drone, Side::B, "Drone.0") => {
                            *m = gltf
                                .named_materials
                                .get("DroneBlue")
                                .expect("Load Blue Drone texture")
                                .clone()
                        }
                        (Piece::Drone, Side::B, "Drone.1") => {
                            *m = gltf
                                .named_materials
                                .get("DotsBlue")
                                .expect("Load Blue Dots texture")
                                .clone()
                        }
                        (Piece::Pawn, Side::A, "Pawn.0") => {
                            *m = gltf
                                .named_materials
                                .get("Pawn")
                                .expect("Load Red Pawn texture")
                                .clone()
                        }
                        (Piece::Pawn, Side::A, "Pawn.1") => {
                            *m = gltf
                                .named_materials
                                .get("Dots")
                                .expect("Load Red Dots texture")
                                .clone()
                        }
                        (Piece::Pawn, Side::B, "Pawn.0") => {
                            *m = gltf
                                .named_materials
                                .get("DroneBlue") // TODO: FIX
                                .expect("Load Blue Pawn texture")
                                .clone()
                        }
                        (Piece::Pawn, Side::B, "Pawn.1") => {
                            *m = gltf
                                .named_materials
                                .get("DotsBlue")
                                .expect("Load Blue Dots texture")
                                .clone()
                        }
                        _ => warn!("???"),
                    }
                }
            });
        }
    })
}

/// Function for selecting entities based on ray intersection
/// There is a bug where we are selecting multiple entities...
fn select_3d(
    mut events: EventReader<MouseButtonInput>,
    query: Query<(Entity, &Handle<Mesh>, &GlobalTransform)>,
    meshes: Res<Assets<Mesh>>,
    cameras: Query<(&Camera, &GlobalTransform)>,
    windows: Query<&Window, With<PrimaryWindow>>,
    parents: Query<Entity, (With<game::Piece>, With<Display3d>)>,
    children: Query<&Children>,
    mut commands: Commands,
    selected: Query<Entity, (With<game::Selected>, With<game::Piece>, With<Display3d>)>,
) {
    events
        .iter()
        .filter(|ev| ev.state == ButtonState::Pressed)
        .for_each(|_| {
            windows.iter().for_each(|window| {
                window.cursor_position().and_then(|pos| {
                    cameras.iter().for_each(|(camera, gt)| {
                        camera.viewport_to_world(gt, pos).and_then(|ray| {
                            query
                                .iter()
                                .filter_map(|(entity, handle, gt)| {
                                    meshes.get(handle).map(|mesh| (entity, mesh, gt))
                                })
                                .for_each(|(entity, mesh, gt)| {
                                    hit3d::intersects(&ray, mesh, &gt).and_then(|_hit| {
                                        parents
                                            .iter()
                                            .find(|&parent| {
                                                children
                                                    .iter_descendants(parent)
                                                    .any(|child| child == entity)
                                            })
                                            .iter()
                                            .for_each(|&parent| {
                                                selected.iter().for_each(|s| {
                                                    if s != parent {
                                                        commands
                                                            .entity(s)
                                                            .remove::<game::Selected>();
                                                    }
                                                });
                                                commands.entity(parent).insert(game::Selected);
                                            });
                                        Some(())
                                    });
                                });
                            Some(())
                        });
                    });
                    Some(())
                });
            });
        });
}

fn selected_gizmo(
    selected: Query<&Transform, (With<game::Selected>, With<game::Piece>, With<Display3d>)>,
    mut gizmos: Gizmos,
) {
    selected.iter().for_each(|transform| {
        gizmos.cuboid(transform.clone(), Color::GREEN);
    })
}

fn pick_up(
    mut events: Query<
        (Entity, &game::Piece),
        (With<game::Piece>, With<Display3d>, Added<game::Selected>),
    >,
    assets_map: Res<AssetsMap>,
    gltfs: Res<Assets<Gltf>>,
    children: Query<&Children>,
    mut players: Query<&mut AnimationPlayer>,
) {
    events.iter_mut().for_each(|(entity, piece)| {
        let gltf = gltfs.get(&assets_map.models).expect("Load GLTF content");
        children.iter_descendants(entity).for_each(|child| {
            if let Ok(mut player) = players.get_mut(child) {
                let animation = match piece {
                    game::Piece::Queen => gltf.named_animations.get("QueenPickup"),
                    game::Piece::Drone => gltf.named_animations.get("DronePickup"),
                    game::Piece::Pawn => gltf.named_animations.get("PawnPickup"),
                };
                player.play(animation.expect("Pickup Animation").clone());
            }
        })
    });
}

fn put_down(
    mut events: RemovedComponents<game::Selected>,
    mut query: Query<&game::Piece, (With<game::Piece>, With<Display3d>)>,
    assets_map: Res<AssetsMap>,
    gltfs: Res<Assets<Gltf>>,
    children: Query<&Children>,
    mut players: Query<&mut AnimationPlayer>,
) {
    events.iter().for_each(|entity| {
        if let Ok(piece) = query.get_mut(entity) {
            let gltf = gltfs.get(&assets_map.models).expect("Load GLTF content");
            children.iter_descendants(entity).for_each(|child| {
                if let Ok(mut player) = players.get_mut(child) {
                    let animation = match piece {
                        game::Piece::Queen => gltf.named_animations.get("QueenPutDown"),
                        game::Piece::Drone => gltf.named_animations.get("DronePutDown"),
                        game::Piece::Pawn => gltf.named_animations.get("PawnPutDown"),
                    };
                    player.play(animation.expect("PutDown Animation").clone());
                }
            })
        }
    })
}

// Animations
// * QueenPickup
// * QueenPutDown
// * QueenIdle
// * DronePickup
// * DronePutDown
// * DroneIdle
// * PawnPickup
// * PawnPutDown
// * PawnIdle