Blacksite/crates/editor/src/viewport/viewport_mode.rs

//! Viewport shading modes: Lit, Unlit (base color), and Collider-only.

use bevy::prelude::*;

use crate::infra::EditorOnly;
use crate::selection_outline::SelectionOutlineShell;
use crate::state::scene_tools_active;
use crate::visualizers::EditorVisualizerProxy;

/// Viewport display mode (editor-only; does not affect saved scenes or runtime rendering).
#[derive(Resource, Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum EditorViewportMode {
    #[default]
    Lit,
    /// Base-color albedo without lighting response (debug).
    Unlit,
    /// Hide scene meshes; emphasize collider gizmos.
    Collider,
}

#[derive(Component)]
struct ViewportMaterialOverride {
    original: Handle<StandardMaterial>,
}

#[derive(Component)]
struct ViewportMeshHiddenByMode;

pub struct EditorViewportModePlugin;

impl Plugin for EditorViewportModePlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<EditorViewportMode>().add_systems(
            Update,
            (
                apply_viewport_mode_on_change,
                apply_viewport_mode_for_new_meshes,
                reconcile_lit_material_overrides,
                reconcile_lit_mode_hidden_meshes,
            )
                .chain()
                .run_if(scene_tools_active),
        );
    }
}

fn apply_viewport_mode_on_change(
    mode: Res<EditorViewportMode>,
    mut last: Local<Option<EditorViewportMode>>,
    mut commands: Commands,
    mut materials: ResMut<Assets<StandardMaterial>>,
    meshes: Query<
        Entity,
        (
            With<Mesh3d>,
            Without<EditorOnly>,
            Without<EditorVisualizerProxy>,
            Without<SelectionOutlineShell>,
        ),
    >,
    mesh_materials: Query<(Entity, &MeshMaterial3d<StandardMaterial>)>,
    overrides: Query<(Entity, &ViewportMaterialOverride)>,
    hidden: Query<Entity, With<ViewportMeshHiddenByMode>>,
    mut visibility: Query<&mut Visibility>,
) {
    let mode = *mode;
    if *last == Some(mode) {
        return;
    }
    *last = Some(mode);

    restore_lit_state(&mut commands, &overrides, &hidden, &mut visibility);

    match mode {
        EditorViewportMode::Lit => {}
        EditorViewportMode::Unlit => {
            for (entity, material) in &mesh_materials {
                apply_unlit_override(entity, material, &mut commands, &mut materials);
            }
        }
        EditorViewportMode::Collider => {
            for entity in meshes.iter() {
                hide_mesh_for_collider_mode(entity, &mut commands, &mut visibility);
            }
        }
    }
}

fn apply_viewport_mode_for_new_meshes(
    mode: Res<EditorViewportMode>,
    mut commands: Commands,
    mut materials: ResMut<Assets<StandardMaterial>>,
    refreshed_materials: Query<
        (Entity, &MeshMaterial3d<StandardMaterial>),
        (
            Changed<MeshMaterial3d<StandardMaterial>>,
            With<ViewportMaterialOverride>,
        ),
    >,
    added_meshes: Query<
        Entity,
        (
            Added<Mesh3d>,
            Without<EditorOnly>,
            Without<EditorVisualizerProxy>,
            Without<SelectionOutlineShell>,
            Without<ViewportMeshHiddenByMode>,
        ),
    >,
    added_materials: Query<
        (Entity, &MeshMaterial3d<StandardMaterial>),
        (
            Or<(
                Added<MeshMaterial3d<StandardMaterial>>,
                Changed<MeshMaterial3d<StandardMaterial>>,
            )>,
            Without<ViewportMaterialOverride>,
        ),
    >,
    mut visibility: Query<&mut Visibility>,
) {
    match *mode {
        EditorViewportMode::Lit => {}
        EditorViewportMode::Unlit => {
            for (entity, material) in &refreshed_materials {
                commands.entity(entity).remove::<ViewportMaterialOverride>();
                apply_unlit_override(entity, material, &mut commands, &mut materials);
            }
            for (entity, material) in &added_materials {
                apply_unlit_override(entity, material, &mut commands, &mut materials);
            }
        }
        EditorViewportMode::Collider => {
            for entity in added_meshes.iter() {
                hide_mesh_for_collider_mode(entity, &mut commands, &mut visibility);
            }
        }
    }
}

fn restore_lit_state(
    commands: &mut Commands,
    overrides: &Query<(Entity, &ViewportMaterialOverride)>,
    hidden: &Query<Entity, With<ViewportMeshHiddenByMode>>,
    visibility: &mut Query<&mut Visibility>,
) {
    for (entity, backup) in overrides.iter() {
        commands
            .entity(entity)
            .insert(MeshMaterial3d(backup.original.clone()));
        commands.entity(entity).remove::<ViewportMaterialOverride>();
    }

    for entity in hidden.iter() {
        if let Ok(mut vis) = visibility.get_mut(entity) {
            *vis = Visibility::Inherited;
        }
        commands.entity(entity).remove::<ViewportMeshHiddenByMode>();
    }
}

fn apply_unlit_override(
    entity: Entity,
    material: &MeshMaterial3d<StandardMaterial>,
    commands: &mut Commands,
    materials: &mut Assets<StandardMaterial>,
) {
    let Some(source) = materials.get(&material.0) else {
        return;
    };
    let unlit = materials.add(StandardMaterial {
        base_color: source.base_color,
        base_color_texture: source.base_color_texture.clone(),
        unlit: true,
        metallic: 0.0,
        perceptual_roughness: 1.0,
        ..default()
    });
    let original = material.0.clone();
    commands
        .entity(entity)
        .insert((ViewportMaterialOverride { original }, MeshMaterial3d(unlit)));
}

fn hide_mesh_for_collider_mode(
    entity: Entity,
    commands: &mut Commands,
    visibility: &mut Query<&mut Visibility>,
) {
    if let Ok(mut vis) = visibility.get_mut(entity) {
        *vis = Visibility::Hidden;
    }
    commands.entity(entity).insert(ViewportMeshHiddenByMode);
}

/// Drops stale unlit overrides after material re-hydration while Lit is active.
fn reconcile_lit_material_overrides(
    mode: Res<EditorViewportMode>,
    mut commands: Commands,
    stale: Query<
        Entity,
        (
            With<ViewportMaterialOverride>,
            Changed<MeshMaterial3d<StandardMaterial>>,
        ),
    >,
) {
    if *mode != EditorViewportMode::Lit {
        return;
    }
    for entity in &stale {
        commands.entity(entity).remove::<ViewportMaterialOverride>();
    }
}

/// Clears stale Collider-mode hides while Lit is active (e.g. after hot reload).
fn reconcile_lit_mode_hidden_meshes(
    mode: Res<EditorViewportMode>,
    mut commands: Commands,
    overrides: Query<(Entity, &ViewportMaterialOverride)>,
    hidden: Query<Entity, With<ViewportMeshHiddenByMode>>,
    mut visibility: Query<&mut Visibility>,
) {
    if *mode != EditorViewportMode::Lit {
        return;
    }
    if overrides.is_empty() && hidden.is_empty() {
        return;
    }
    restore_lit_state(&mut commands, &overrides, &hidden, &mut visibility);
}