From 6c34ab8bd64708194a2945669634ec1219bf72d0 Mon Sep 17 00:00:00 2001
From: Jan Hohenheim <jan@hohenheim.ch>
Date: Mon, 25 Mar 2024 11:32:04 +0100
Subject: [PATCH] feat(bevy_gltf_components): Improve global illumination to
 match Blender ambient lighting better(#174)

* Improves global illumination/ ambient by switching to an env map
---
 .../src/blender_settings/lighting.rs          | 69 +++++++++++++++++--
 1 file changed, 63 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_gltf_components/src/blender_settings/lighting.rs b/crates/bevy_gltf_components/src/blender_settings/lighting.rs
index 622c0e1..6d758a4 100644
--- a/crates/bevy_gltf_components/src/blender_settings/lighting.rs
+++ b/crates/bevy_gltf_components/src/blender_settings/lighting.rs
@@ -1,5 +1,10 @@
 use bevy::pbr::DirectionalLightShadowMap;
 use bevy::prelude::*;
+use bevy::render::render_asset::RenderAssetUsages;
+use bevy::render::render_resource::{
+    Extent3d, TextureDimension, TextureFormat, TextureViewDescriptor, TextureViewDimension,
+};
+use std::iter;
 
 use crate::GltfComponentsSet;
 
@@ -84,14 +89,66 @@ fn process_shadowmap(
 }
 
 fn process_background_shader(
-    background_shaders: Query<&BlenderBackgroundShader, Added<BlenderBackgroundShader>>,
+    background_shaders: Query<Ref<BlenderBackgroundShader>>,
+    cameras: Query<(Entity, Ref<Camera3d>)>,
+    mut images: ResMut<Assets<Image>>,
     mut commands: Commands,
+    mut env_map_handle: Local<Option<Handle<Image>>>,
 ) {
-    for background_shader in background_shaders.iter() {
-        commands.insert_resource(AmbientLight {
-            color: background_shader.color,
-            // Just a guess, see <https://github.com/bevyengine/bevy/issues/12280>
-            brightness: background_shader.strength * 400.0,
+    let Ok(background_shader) = background_shaders.get_single() else {
+        return;
+    };
+
+    let env_map_handle = env_map_handle.get_or_insert_with(|| {
+        let size = Extent3d {
+            width: 1,
+            height: 6,
+            depth_or_array_layers: 1,
+        };
+        let dimension = TextureDimension::D2;
+        const SIDES_PER_CUBE: usize = 6;
+        let data: Vec<_> = iter::repeat(background_shader.color.as_rgba_u8())
+            .take(SIDES_PER_CUBE)
+            .flatten()
+            .collect();
+        let format = TextureFormat::Rgba8UnormSrgb;
+        let asset_usage = RenderAssetUsages::RENDER_WORLD;
+
+        let mut image = Image::new(size, dimension, data, format, asset_usage);
+
+        // Source: https://github.com/bevyengine/bevy/blob/85b488b73d6f6e75690962fba67a144d9beb6b88/examples/3d/skybox.rs#L152-L160
+        image.reinterpret_stacked_2d_as_array(image.height() / image.width());
+        image.texture_view_descriptor = Some(TextureViewDescriptor {
+            dimension: Some(TextureViewDimension::Cube),
+            ..default()
+        });
+
+        images.add(image)
+    });
+    // Don't need the handle to be &mut
+    let env_map_handle = &*env_map_handle;
+
+    if background_shader.is_added() {
+        // We're using an environment map, so we don't need the ambient light
+        commands.remove_resource::<AmbientLight>();
+    }
+
+    let is_bg_outdated = background_shader.is_changed();
+    if is_bg_outdated {
+        let color = background_shader.color * background_shader.strength;
+        commands.insert_resource(ClearColor(color));
+    }
+    let camera_entities = cameras
+        .iter()
+        .filter_map(|(entity, cam)| (is_bg_outdated || cam.is_changed()).then_some(entity));
+
+    for camera_entity in camera_entities {
+        // See https://github.com/KhronosGroup/glTF-Blender-IO/blob/8573cc0dfb612091bfc1bcf6df55c18a44b9668a/addons/io_scene_gltf2/blender/com/gltf2_blender_conversion.py#L19
+        const PBR_WATTS_TO_LUMENS: f32 = 683.0;
+        commands.entity(camera_entity).insert(EnvironmentMapLight {
+            diffuse_map: env_map_handle.clone(),
+            specular_map: env_map_handle.clone(),
+            intensity: background_shader.strength * PBR_WATTS_TO_LUMENS,
         });
     }
 }