Blacksite/crates/editor/src/assets/thumbnails/studio.rs

//! Offscreen model thumbnail studio for FBX and untextured glTF assets.

use std::collections::VecDeque;
use std::path::Path;

use bevy::camera::visibility::RenderLayers;
use bevy::camera::RenderTarget;
use bevy::gltf::GltfAssetLabel;
use bevy::math::bounding::Aabb3d;
use bevy::mesh::VertexAttributeValues;
use bevy::prelude::*;
use bevy::render::render_resource::TextureFormat;
use bevy_egui::EguiUserTextures;
use shared::{material_from_desc, standard_material_asset_path, ModelRef};

use super::cache::AssetThumbnailCache;
use super::job::{ThumbnailJob, ThumbnailJobSource};
use super::sources::{source_for_extension, uses_scene_root};
use crate::assets::asset_server_path;
use crate::infra::EditorOnly;

pub(crate) const THUMB_SIZE: u32 = 128;
pub(crate) const THUMBNAIL_LAYER: usize = 31;
const THUMBNAIL_VERTICAL_FOV: f32 = 40.0_f32.to_radians();
const THUMBNAIL_HALF_VERTICAL_FOV: f32 = THUMBNAIL_VERTICAL_FOV * 0.5;
const SCENE_FRAME_PADDING: f32 = 1.35;
const MATERIAL_SPHERE_FRAME_PADDING: f32 = 1.28;
const MESH_WARMUP_FRAMES: u8 = 12;
const POST_ATTACH_FRAMES: u8 = 8;
const RENDER_FRAMES: u8 = 3;
const COOLDOWN_FRAMES: u8 = 2;
const LOAD_TIMEOUT_FRAMES: u32 = 240;

#[derive(Component)]
struct ThumbnailStudioCamera;

#[derive(Component)]
struct ThumbnailStudioLight;

#[derive(Component)]
struct ThumbnailStudioLayer;

#[derive(Component)]
struct ThumbnailSceneRoot;

#[derive(Resource)]
pub struct ThumbnailStudio {
    camera: Entity,
    render_image: Handle<Image>,
    _lights: Vec<Entity>,
    queue: VecDeque<ThumbnailJob>,
    active: Option<ActiveModelThumbnail>,
    cooldown_frames: u8,
}

struct ActiveModelThumbnail {
    cache_key: String,
    label: String,
    render_image: Handle<Image>,
    root: Entity,
    meshes_ready: bool,
    mesh_warmup_frames: u8,
    camera_active: bool,
    post_attach_frames: u8,
    framed: bool,
    frames_remaining: u8,
    wait_frames: u32,
    framing: ThumbnailFraming,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ThumbnailFraming {
    SceneBounds,
    MaterialSphere,
}

impl ThumbnailFraming {
    fn for_source(source: &ThumbnailJobSource) -> Self {
        match source {
            ThumbnailJobSource::SourceMaterial { .. }
            | ThumbnailJobSource::MaterialAsset { .. } => Self::MaterialSphere,
            ThumbnailJobSource::Model { .. } | ThumbnailJobSource::MeshSubAsset { .. } => {
                Self::SceneBounds
            }
        }
    }
}

pub struct ThumbnailStudioPlugin;

impl Plugin for ThumbnailStudioPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Startup, setup_thumbnail_studio)
            .add_systems(PostUpdate, apply_thumbnail_studio_layers)
            .add_systems(Last, process_thumbnail_studio);
    }
}

impl ThumbnailStudio {
    pub fn enqueue(&mut self, cache_key: String, model_path: String) -> bool {
        self.enqueue_source(cache_key, ThumbnailJobSource::Model { model_path })
    }

    pub fn enqueue_source(&mut self, cache_key: String, source: ThumbnailJobSource) -> bool {
        if self
            .active
            .as_ref()
            .is_some_and(|active| active.cache_key == cache_key)
        {
            return false;
        }
        if self.queue.iter().any(|job| job.cache_key == cache_key) {
            return false;
        }
        self.queue.push_back(ThumbnailJob { cache_key, source });
        true
    }

    pub fn clear_queue(&mut self) {
        self.queue.clear();
        self.cooldown_frames = 0;
    }

    pub fn clear_jobs(&mut self) -> StudioCleanup {
        StudioCleanup {
            camera: self.camera,
            active_root: self.active.take().map(|active| active.root),
        }
    }
}

pub struct StudioCleanup {
    pub camera: Entity,
    pub active_root: Option<Entity>,
}

pub fn model_scene_asset_path(path: &str, scene_index: usize) -> String {
    let path = asset_server_path(path);
    if path.ends_with(".fbx") {
        ModelRef::fbx_scene_asset_path(&path, scene_index)
    } else {
        GltfAssetLabel::Scene(scene_index)
            .from_asset(path)
            .to_string()
    }
}

fn setup_thumbnail_studio(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
    let render_image = images.add(studio_render_image());

    let camera = commands
        .spawn((
            Name::new("Thumbnail Studio Camera"),
            EditorOnly,
            ThumbnailStudioCamera,
            ThumbnailStudioLayer,
            RenderLayers::layer(THUMBNAIL_LAYER),
            Camera3d::default(),
            Msaa::Off,
            Camera {
                is_active: false,
                order: -50,
                clear_color: ClearColorConfig::Custom(Color::srgb(0.12, 0.12, 0.14)),
                ..default()
            },
            RenderTarget::Image(render_image.clone().into()),
            Transform::from_xyz(2.0, 1.4, 2.0).looking_at(Vec3::new(0.0, 0.5, 0.0), Vec3::Y),
            Projection::Perspective(PerspectiveProjection {
                fov: THUMBNAIL_VERTICAL_FOV,
                ..default()
            }),
        ))
        .id();

    let key_light = commands
        .spawn((
            EditorOnly,
            ThumbnailStudioLight,
            ThumbnailStudioLayer,
            RenderLayers::layer(THUMBNAIL_LAYER),
            DirectionalLight {
                illuminance: 12_000.0,
                shadow_maps_enabled: false,
                ..default()
            },
            Transform::from_rotation(Quat::from_euler(EulerRot::XYZ, -0.8, 0.9, 0.0)),
        ))
        .id();

    let fill_light = commands
        .spawn((
            EditorOnly,
            ThumbnailStudioLight,
            ThumbnailStudioLayer,
            RenderLayers::layer(THUMBNAIL_LAYER),
            DirectionalLight {
                illuminance: 3_500.0,
                shadow_maps_enabled: false,
                ..default()
            },
            Transform::from_rotation(Quat::from_euler(EulerRot::XYZ, -0.4, -2.2, 0.0)),
        ))
        .id();

    commands.insert_resource(ThumbnailStudio {
        camera,
        render_image,
        _lights: vec![key_light, fill_light],
        queue: VecDeque::new(),
        active: None,
        cooldown_frames: 0,
    });
}

#[allow(clippy::too_many_arguments)]
fn process_thumbnail_studio(
    mut commands: Commands,
    mut studio: ResMut<ThumbnailStudio>,
    mut cache: ResMut<AssetThumbnailCache>,
    mut textures: ResMut<EguiUserTextures>,
    mut materials: ResMut<Assets<StandardMaterial>>,
    mut mesh_storage: ResMut<Assets<Mesh>>,
    mut images: ResMut<Assets<Image>>,
    asset_server: Res<AssetServer>,
    mut cameras: Query<&mut Camera, With<ThumbnailStudioCamera>>,
    mut camera_targets: Query<&mut RenderTarget, With<ThumbnailStudioCamera>>,
    mut camera_transforms: Query<&mut Transform, With<ThumbnailStudioCamera>>,
    children: Query<&Children>,
    transforms: Query<&GlobalTransform>,
    mesh3d: Query<&Mesh3d>,
) {
    if studio.cooldown_frames > 0 {
        studio.cooldown_frames -= 1;
        return;
    }

    if let Some(mut active) = studio.active.take() {
        if !active.framed {
            active.wait_frames += 1;

            if !active.meshes_ready {
                if thumbnail_meshes_present(active.root, &children, &mesh3d, &mesh_storage) {
                    active.meshes_ready = true;
                    active.mesh_warmup_frames = 0;
                }
            } else if active.mesh_warmup_frames < MESH_WARMUP_FRAMES {
                active.mesh_warmup_frames += 1;
            } else if frame_thumbnail_scene(
                active.root,
                &children,
                &transforms,
                &mesh3d,
                &mesh_storage,
                &mut camera_transforms,
                active.framing,
            ) {
                active.framed = true;
                active.post_attach_frames = 0;
                studio.active = Some(active);
                return;
            }

            if active.wait_frames >= LOAD_TIMEOUT_FRAMES {
                warn!(
                    "Model thumbnail timed out framing {} (meshes may still be loading)",
                    active.label
                );
                fail_active_thumbnail(
                    &mut commands,
                    &mut studio,
                    &mut cache,
                    &mut cameras,
                    &mut camera_targets,
                    active,
                    "load timeout",
                    true,
                );
            } else {
                studio.active = Some(active);
            }
            return;
        }

        if !active.camera_active {
            active.post_attach_frames += 1;
            if active.post_attach_frames >= POST_ATTACH_FRAMES {
                if let Ok(mut camera) = cameras.get_mut(studio.camera) {
                    camera.is_active = true;
                }
                active.camera_active = true;
                active.frames_remaining = RENDER_FRAMES;
            }
            studio.active = Some(active);
            return;
        }

        if active.camera_active {
            if active.frames_remaining > 0 {
                active.frames_remaining -= 1;
                studio.active = Some(active);
            } else {
                cache.complete_studio_thumbnail(
                    &active.cache_key,
                    active.render_image.clone(),
                    &mut textures,
                );
                finish_active_thumbnail(
                    &mut commands,
                    &mut studio,
                    &mut cameras,
                    &mut camera_targets,
                    active,
                );
            }
            return;
        }

        studio.active = Some(active);
        return;
    }

    let Some(job) = studio.queue.pop_front() else {
        deactivate_studio_camera(&mut cameras, studio.camera);
        return;
    };

    let render_image = images.add(studio_render_image());
    if let Ok(mut target) = camera_targets.get_mut(studio.camera) {
        *target = RenderTarget::Image(render_image.clone().into());
    }
    if let Ok(mut camera) = cameras.get_mut(studio.camera) {
        camera.is_active = false;
    }

    let root = commands
        .spawn((
            EditorOnly,
            ThumbnailSceneRoot,
            ThumbnailStudioLayer,
            RenderLayers::layer(THUMBNAIL_LAYER),
            Transform::default(),
            Visibility::Visible,
            InheritedVisibility::default(),
            ViewVisibility::default(),
        ))
        .id();

    let label = job.source.label().to_string();
    let framing = ThumbnailFraming::for_source(&job.source);
    let content_spawned = spawn_thumbnail_job_content(
        &mut commands,
        &mut mesh_storage,
        &mut materials,
        &asset_server,
        root,
        &job.source,
    );

    if !content_spawned {
        fail_active_thumbnail(
            &mut commands,
            &mut studio,
            &mut cache,
            &mut cameras,
            &mut camera_targets,
            ActiveModelThumbnail {
                cache_key: job.cache_key,
                label,
                render_image,
                root,
                meshes_ready: false,
                mesh_warmup_frames: 0,
                camera_active: false,
                post_attach_frames: 0,
                framed: false,
                frames_remaining: 0,
                wait_frames: 0,
                framing,
            },
            "no renderable meshes",
            true,
        );
        return;
    }

    studio.active = Some(ActiveModelThumbnail {
        cache_key: job.cache_key,
        label,
        render_image,
        root,
        meshes_ready: false,
        mesh_warmup_frames: 0,
        camera_active: false,
        post_attach_frames: 0,
        framed: false,
        frames_remaining: 0,
        wait_frames: 0,
        framing,
    });
}

fn spawn_thumbnail_job_content(
    commands: &mut Commands,
    mesh_storage: &mut Assets<Mesh>,
    materials: &mut Assets<StandardMaterial>,
    asset_server: &AssetServer,
    root: Entity,
    source: &ThumbnailJobSource,
) -> bool {
    match source {
        ThumbnailJobSource::Model { model_path } => {
            if uses_scene_root(model_path) {
                commands.entity(root).insert(WorldAssetRoot(
                    asset_server.load(model_scene_asset_path(model_path, 0)),
                ));
            }
            if let Some(source) = source_for_extension(model_path) {
                source.spawn_preview(commands, mesh_storage, materials, root, model_path)
            } else {
                uses_scene_root(model_path)
            }
        }
        ThumbnailJobSource::MeshSubAsset {
            model_path,
            mesh_label,
            material_label,
        } => spawn_mesh_subasset_preview(
            commands,
            materials,
            asset_server,
            root,
            model_path,
            mesh_label,
            material_label.as_deref(),
        ),
        ThumbnailJobSource::SourceMaterial {
            model_path,
            material_label,
        } => spawn_material_sphere(
            commands,
            mesh_storage,
            root,
            asset_server.load(standard_material_asset_path(
                &asset_server_path(model_path),
                material_label,
            )),
        ),
        ThumbnailJobSource::MaterialAsset { material, .. } => {
            let material = materials.add(material_from_desc(asset_server, material));
            spawn_material_sphere(commands, mesh_storage, root, material)
        }
    }
}

fn spawn_mesh_subasset_preview(
    commands: &mut Commands,
    materials: &mut Assets<StandardMaterial>,
    asset_server: &AssetServer,
    root: Entity,
    model_path: &str,
    mesh_label: &str,
    material_label: Option<&str>,
) -> bool {
    let source_path = asset_server_path(model_path);
    let mesh: Handle<Mesh> = asset_server.load(labeled_asset_path(&source_path, mesh_label));
    let material = if let Some(material_label) = material_label {
        asset_server.load(standard_material_asset_path(&source_path, material_label))
    } else {
        materials.add(StandardMaterial {
            base_color: Color::srgb(0.72, 0.72, 0.75),
            perceptual_roughness: 0.65,
            ..default()
        })
    };
    commands.entity(root).with_children(|parent| {
        parent.spawn((
            RenderLayers::layer(THUMBNAIL_LAYER),
            Mesh3d(mesh),
            MeshMaterial3d(material),
            Transform::default(),
        ));
    });
    true
}

fn spawn_material_sphere(
    commands: &mut Commands,
    mesh_storage: &mut Assets<Mesh>,
    root: Entity,
    material: Handle<StandardMaterial>,
) -> bool {
    let mesh = mesh_storage.add(Sphere::new(0.65).mesh().uv(48, 24));
    commands.entity(root).with_children(|parent| {
        parent.spawn((
            RenderLayers::layer(THUMBNAIL_LAYER),
            Mesh3d(mesh),
            MeshMaterial3d(material),
            Transform::default(),
        ));
    });
    true
}

fn labeled_asset_path(path: &str, label: &str) -> String {
    format!("{path}#{label}")
}

fn studio_render_image() -> Image {
    Image::new_target_texture(THUMB_SIZE, THUMB_SIZE, TextureFormat::Rgba16Float, None)
}

fn finish_active_thumbnail(
    commands: &mut Commands,
    studio: &mut ThumbnailStudio,
    cameras: &mut Query<&mut Camera, With<ThumbnailStudioCamera>>,
    camera_targets: &mut Query<&mut RenderTarget, With<ThumbnailStudioCamera>>,
    active: ActiveModelThumbnail,
) {
    despawn_thumbnail_root(commands, active.root);
    if let Ok(mut target) = camera_targets.get_mut(studio.camera) {
        *target = RenderTarget::Image(studio.render_image.clone().into());
    }
    deactivate_studio_camera(cameras, studio.camera);
    studio.cooldown_frames = COOLDOWN_FRAMES;
}

#[expect(
    clippy::too_many_arguments,
    reason = "thumbnail failure coordinates the job, cache, render target, and camera state"
)]
fn fail_active_thumbnail(
    commands: &mut Commands,
    studio: &mut ThumbnailStudio,
    cache: &mut AssetThumbnailCache,
    cameras: &mut Query<&mut Camera, With<ThumbnailStudioCamera>>,
    camera_targets: &mut Query<&mut RenderTarget, With<ThumbnailStudioCamera>>,
    active: ActiveModelThumbnail,
    reason: &str,
    retryable: bool,
) {
    cache.mark_studio_failed(&active.cache_key, reason, retryable);
    finish_active_thumbnail(commands, studio, cameras, camera_targets, active);
}

fn apply_thumbnail_studio_layers(
    mut commands: Commands,
    roots: Query<Entity, With<ThumbnailSceneRoot>>,
    children: Query<&Children>,
    tagged: Query<(), With<ThumbnailStudioLayer>>,
) {
    for root in &roots {
        let mut stack = vec![root];
        while let Some(entity) = stack.pop() {
            if tagged.get(entity).is_err() {
                commands
                    .entity(entity)
                    .insert((ThumbnailStudioLayer, RenderLayers::layer(THUMBNAIL_LAYER)));
            }
            if let Ok(kids) = children.get(entity) {
                stack.extend(kids.iter());
            }
        }
    }
}

fn deactivate_studio_camera(
    cameras: &mut Query<&mut Camera, With<ThumbnailStudioCamera>>,
    camera: Entity,
) {
    if let Ok(mut camera) = cameras.get_mut(camera) {
        camera.is_active = false;
    }
}

fn despawn_thumbnail_root(commands: &mut Commands, root: Entity) {
    commands.entity(root).despawn_children();
    commands.entity(root).despawn();
}

fn thumbnail_meshes_present(
    root: Entity,
    children: &Query<&Children>,
    meshes: &Query<&Mesh3d>,
    mesh_assets: &Assets<Mesh>,
) -> bool {
    let mut entities = Vec::new();
    collect_descendants(root, children, &mut entities);

    entities.iter().any(|&entity| {
        let Ok(mesh3d) = meshes.get(entity) else {
            return false;
        };
        mesh_assets
            .get(&mesh3d.0)
            .is_some_and(|mesh| mesh_positions(mesh).next().is_some())
    })
}

fn mesh_positions(mesh: &Mesh) -> impl Iterator<Item = [f32; 3]> + '_ {
    let positions = mesh
        .attribute(Mesh::ATTRIBUTE_POSITION)
        .and_then(|values| match values {
            VertexAttributeValues::Float32x3(positions) => Some(positions.as_slice()),
            _ => None,
        });
    positions
        .into_iter()
        .flat_map(|slice| slice.iter().copied())
}

fn frame_thumbnail_scene(
    root: Entity,
    children: &Query<&Children>,
    transforms: &Query<&GlobalTransform>,
    meshes: &Query<&Mesh3d>,
    mesh_assets: &Assets<Mesh>,
    camera_transforms: &mut Query<&mut Transform, With<ThumbnailStudioCamera>>,
    framing: ThumbnailFraming,
) -> bool {
    let mut entities = Vec::new();
    collect_descendants(root, children, &mut entities);

    let mut min = Vec3::splat(f32::INFINITY);
    let mut max = Vec3::splat(f32::NEG_INFINITY);
    let mut found_mesh = false;

    for entity in entities {
        let Ok(mesh3d) = meshes.get(entity) else {
            continue;
        };
        let Ok(global) = transforms.get(entity) else {
            continue;
        };

        if let Some(mesh) = mesh_assets.get(&mesh3d.0) {
            let mut local_min = Vec3::splat(f32::INFINITY);
            let mut local_max = Vec3::splat(f32::NEG_INFINITY);
            let mut has_positions = false;
            for position in mesh_positions(mesh) {
                let p = Vec3::from_array(position);
                local_min = local_min.min(p);
                local_max = local_max.max(p);
                has_positions = true;
            }
            if has_positions {
                found_mesh = true;
                for corner in aabb_corners(&Aabb3d::from_min_max(
                    Vec3A::from(local_min),
                    Vec3A::from(local_max),
                )) {
                    let world = global.transform_point(corner);
                    min = min.min(world);
                    max = max.max(world);
                }
            }
        }
    }

    if !found_mesh {
        return false;
    }

    let center = (min + max) * 0.5;
    let distance = thumbnail_camera_distance(min, max, framing);
    let eye = center + Vec3::new(1.0, 0.75, 1.0).normalize() * distance;

    if let Ok(mut transform) = camera_transforms.single_mut() {
        *transform = Transform::from_translation(eye).looking_at(center, Vec3::Y);
    }

    true
}

fn thumbnail_camera_distance(min: Vec3, max: Vec3, framing: ThumbnailFraming) -> f32 {
    let bounds_size = max - min;
    let (radius, padding) = match framing {
        ThumbnailFraming::SceneBounds => {
            (bounds_size.length().max(0.25) * 0.5, SCENE_FRAME_PADDING)
        }
        ThumbnailFraming::MaterialSphere => (
            bounds_size.max_element().max(0.5) * 0.5,
            MATERIAL_SPHERE_FRAME_PADDING,
        ),
    };
    radius / THUMBNAIL_HALF_VERTICAL_FOV.tan() * padding
}

fn collect_descendants(entity: Entity, children: &Query<&Children>, out: &mut Vec<Entity>) {
    out.push(entity);
    let Ok(child_list) = children.get(entity) else {
        return;
    };
    for child in child_list.iter() {
        collect_descendants(child, children, out);
    }
}

fn aabb_corners(aabb: &Aabb3d) -> [Vec3; 8] {
    let min = Vec3::from(aabb.min);
    let max = Vec3::from(aabb.max);
    [
        Vec3::new(min.x, min.y, min.z),
        Vec3::new(max.x, min.y, min.z),
        Vec3::new(min.x, max.y, min.z),
        Vec3::new(max.x, max.y, min.z),
        Vec3::new(min.x, min.y, max.z),
        Vec3::new(max.x, min.y, max.z),
        Vec3::new(min.x, max.y, max.z),
        Vec3::new(max.x, max.y, max.z),
    ]
}

/// Skip studio jobs when the source file is missing on disk.
pub fn model_file_exists(catalog_path: &str) -> bool {
    Path::new(catalog_path).is_file()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn material_sphere_framing_fills_more_of_the_thumbnail_with_padding() {
        let sphere_radius = 0.65;
        let min = Vec3::splat(-sphere_radius);
        let max = Vec3::splat(sphere_radius);

        let scene_distance = thumbnail_camera_distance(min, max, ThumbnailFraming::SceneBounds);
        let material_distance =
            thumbnail_camera_distance(min, max, ThumbnailFraming::MaterialSphere);
        let projected_radius_fraction =
            sphere_radius / (material_distance * THUMBNAIL_HALF_VERTICAL_FOV.tan());

        assert!(material_distance < scene_distance * 0.6);
        assert!((0.75..0.82).contains(&projected_radius_fraction));
    }
}