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

//! Thumbnail cache for asset browser grid cells.

use std::collections::{HashMap, HashSet};

use bevy::prelude::*;
use bevy_egui::{egui, EguiPrimaryContextPass, EguiTextureHandle, EguiUserTextures};
use egui_phosphor_icons::icons;

use super::sources::gltf::{gltf_base_color_texture_path, validate_gltf_dependencies};
use super::studio::{model_file_exists, ThumbnailStudio};
use super::ThumbnailJobSource;
use crate::assets::{
    asset_cache_key, asset_server_path, EditorAsset, EditorAssetKind, EditorAssets,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ThumbnailState {
    Pending,
    Ready,
    Failed { reason: String, retryable: bool },
}

#[derive(Debug, Clone)]
struct ThumbnailFailure {
    reason: String,
    retryable: bool,
}

#[derive(Resource, Default)]
pub struct AssetThumbnailCache {
    pub texture_ids: HashMap<String, egui::TextureId>,
    pending: HashMap<String, Handle<Image>>,
    pub(crate) studio_pending: HashSet<String>,
    failed: HashMap<String, ThumbnailFailure>,
    prefetched_folder: Option<String>,
}

pub struct AssetThumbnailsPlugin;

impl Plugin for AssetThumbnailsPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<AssetThumbnailCache>()
            .add_systems(EguiPrimaryContextPass, register_loaded_thumbnails);
    }
}

impl AssetThumbnailCache {
    pub fn state(&self, key: &str) -> Option<ThumbnailState> {
        if self.texture_ids.contains_key(key) {
            return Some(ThumbnailState::Ready);
        }
        if self.failed.contains_key(key) {
            let failure = self.failed.get(key)?;
            return Some(ThumbnailState::Failed {
                reason: failure.reason.clone(),
                retryable: failure.retryable,
            });
        }
        if self.pending.contains_key(key) || self.studio_pending.contains(key) {
            return Some(ThumbnailState::Pending);
        }
        None
    }

    pub fn retry(&mut self, key: &str) {
        self.failed.remove(key);
        self.pending.remove(key);
        self.studio_pending.remove(key);
        self.texture_ids.remove(key);
    }

    pub fn request_texture(&mut self, key: String, path: String, asset_server: &AssetServer) {
        if matches!(
            self.state(&key),
            Some(ThumbnailState::Ready | ThumbnailState::Pending)
        ) {
            return;
        }
        if self.failed.contains_key(&key) {
            return;
        }
        let handle: Handle<Image> = asset_server.load(asset_server_path(&path));
        self.pending.insert(key, handle);
    }

    pub fn request_model(
        &mut self,
        key: String,
        model_path: String,
        asset_server: &AssetServer,
        studio: &mut ThumbnailStudio,
    ) {
        if matches!(
            self.state(&key),
            Some(ThumbnailState::Ready | ThumbnailState::Pending)
        ) {
            return;
        }
        if self.failed.contains_key(&key) {
            return;
        }

        if !self.model_source_ready(&key, &model_path) {
            return;
        }

        if let Some(texture_path) = gltf_base_color_texture_path(&model_path) {
            self.request_texture(key, texture_path, asset_server);
            return;
        }

        if studio.enqueue(key.clone(), model_path) {
            self.studio_pending.insert(key);
        }
    }

    pub fn request_mesh_subasset(
        &mut self,
        key: String,
        model_path: String,
        mesh_label: String,
        material_label: Option<String>,
        studio: &mut ThumbnailStudio,
    ) {
        if matches!(
            self.state(&key),
            Some(ThumbnailState::Ready | ThumbnailState::Pending)
        ) || self.failed.contains_key(&key)
        {
            return;
        }
        if !self.model_source_ready(&key, &model_path) {
            return;
        }
        if studio.enqueue_source(
            key.clone(),
            ThumbnailJobSource::MeshSubAsset {
                model_path,
                mesh_label,
                material_label,
            },
        ) {
            self.studio_pending.insert(key);
        }
    }

    pub fn request_source_material(
        &mut self,
        key: String,
        model_path: String,
        material_label: String,
        studio: &mut ThumbnailStudio,
    ) {
        if matches!(
            self.state(&key),
            Some(ThumbnailState::Ready | ThumbnailState::Pending)
        ) || self.failed.contains_key(&key)
        {
            return;
        }
        if !self.model_source_ready(&key, &model_path) {
            return;
        }
        if studio.enqueue_source(
            key.clone(),
            ThumbnailJobSource::SourceMaterial {
                model_path,
                material_label,
            },
        ) {
            self.studio_pending.insert(key);
        }
    }

    pub fn request_material_asset(
        &mut self,
        key: String,
        path: String,
        studio: &mut ThumbnailStudio,
    ) {
        if matches!(
            self.state(&key),
            Some(ThumbnailState::Ready | ThumbnailState::Pending)
        ) || self.failed.contains_key(&key)
        {
            return;
        }
        let material = match shared::MaterialAsset::load_from_path(&path) {
            Ok(asset) => asset,
            Err(error) => {
                self.mark_studio_failed(&key, &error, true);
                return;
            }
        };
        if studio.enqueue_source(
            key.clone(),
            ThumbnailJobSource::MaterialAsset {
                label: material.label.clone(),
                material: Box::new(material.material),
            },
        ) {
            self.studio_pending.insert(key);
        }
    }

    fn model_source_ready(&mut self, key: &str, model_path: &str) -> bool {
        if !model_file_exists(model_path) {
            self.mark_studio_failed(key, "file not found", false);
            return false;
        }
        if let Err(reason) = validate_gltf_dependencies(model_path) {
            self.mark_studio_failed(key, &reason, false);
            return false;
        }
        true
    }

    pub(crate) fn complete_studio_thumbnail(
        &mut self,
        key: &str,
        image: Handle<Image>,
        textures: &mut EguiUserTextures,
    ) {
        self.studio_pending.remove(key);
        self.failed.remove(key);
        let texture_id = textures
            .image_id(&image)
            .unwrap_or_else(|| textures.add_image(EguiTextureHandle::Strong(image)));
        self.texture_ids.insert(key.to_string(), texture_id);
    }

    pub(crate) fn mark_studio_failed(&mut self, key: &str, reason: &str, retryable: bool) {
        self.studio_pending.remove(key);
        self.pending.remove(key);
        self.failed.insert(
            key.to_string(),
            ThumbnailFailure {
                reason: reason.to_string(),
                retryable,
            },
        );
    }

    pub fn invalidate_all(&mut self) {
        self.texture_ids.clear();
        self.pending.clear();
        self.studio_pending.clear();
        self.failed.clear();
        self.prefetched_folder = None;
    }

    pub fn snapshot(&self) -> ThumbnailCacheSnapshot {
        ThumbnailCacheSnapshot {
            texture_ids: self.texture_ids.clone(),
            pending_keys: self
                .pending
                .keys()
                .chain(self.studio_pending.iter())
                .cloned()
                .collect(),
            failed_keys: self.failed.keys().cloned().collect(),
        }
    }
}

fn register_loaded_thumbnails(
    asset_server: Res<AssetServer>,
    mut cache: ResMut<AssetThumbnailCache>,
    mut textures: ResMut<EguiUserTextures>,
) {
    let ready: Vec<(String, Handle<Image>)> = cache
        .pending
        .iter()
        .filter(|(_, handle)| asset_server.is_loaded_with_dependencies(*handle))
        .map(|(key, handle)| (key.clone(), handle.clone()))
        .collect();

    for (key, handle) in ready {
        cache.pending.remove(&key);
        cache.failed.remove(&key);
        let texture_id = textures
            .image_id(&handle)
            .unwrap_or_else(|| textures.add_image(EguiTextureHandle::Strong(handle)));
        cache.texture_ids.insert(key, texture_id);
    }
}

pub fn kind_icon(kind: &EditorAssetKind) -> egui_phosphor_icons::Icon {
    match kind {
        EditorAssetKind::Primitive(_) => icons::CUBE,
        EditorAssetKind::Light(_) => icons::LIGHTBULB,
        EditorAssetKind::Model => icons::CUBE_TRANSPARENT,
        EditorAssetKind::Texture => icons::IMAGE,
        EditorAssetKind::Material => icons::PAINT_BRUSH,
        EditorAssetKind::Level => icons::MAP_TRIFOLD,
        EditorAssetKind::Prefab => icons::PACKAGE,
        EditorAssetKind::PostProcessVolume => icons::CAMERA,
        EditorAssetKind::PostProcessEffect => icons::SPARKLE,
        EditorAssetKind::RenderingProfile => icons::SLIDERS,
        EditorAssetKind::ShaderSchema => icons::CODE,
    }
}

pub fn draw_asset_cell_with(
    ui: &mut egui::Ui,
    asset: &EditorAsset,
    texture_id: Option<egui::TextureId>,
    pending: bool,
    failed: Option<&str>,
    selected: bool,
    thumbnail_size: f32,
) -> egui::Response {
    let thumb_size = thumbnail_size.clamp(48.0, 112.0);
    let cell_size = egui::vec2(thumb_size + 20.0, thumb_size + 30.0);
    let (rect, response) = ui.allocate_exact_size(cell_size, egui::Sense::click_and_drag());

    let fill = if selected {
        crate::ui::theme::SELECTION_BG_MUTED
    } else if response.hovered() {
        crate::ui::theme::ELEVATED_BG
    } else {
        crate::ui::theme::WIDGET_BG
    };
    ui.painter().rect(
        rect,
        4.0,
        fill,
        egui::Stroke::new(
            1.0,
            if selected {
                crate::ui::theme::ACCENT_HOVER
            } else {
                crate::ui::theme::BORDER
            },
        ),
        egui::StrokeKind::Inside,
    );

    let thumb_rect = egui::Rect::from_min_size(
        rect.min + egui::vec2(10.0, 8.0),
        egui::vec2(thumb_size, thumb_size),
    );

    if let Some(texture_id) = texture_id {
        ui.painter().image(
            texture_id,
            thumb_rect,
            egui::Rect::from_min_max(egui::pos2(0.0, 0.0), egui::pos2(1.0, 1.0)),
            egui::Color32::WHITE,
        );
    } else if pending {
        ui.painter().text(
            thumb_rect.center(),
            egui::Align2::CENTER_CENTER,
            icons::CIRCLE_NOTCH.as_str(),
            egui::FontId::new(24.0, egui::FontFamily::Name("phosphor-regular".into())),
            crate::ui::theme::TEXT_DIM,
        );
    } else if failed.is_some() {
        ui.painter().text(
            thumb_rect.center(),
            egui::Align2::CENTER_CENTER,
            icons::WARNING_CIRCLE.as_str(),
            egui::FontId::new(24.0, egui::FontFamily::Name("phosphor-regular".into())),
            crate::ui::theme::TEXT_DIM,
        );
    } else {
        let icon = kind_icon(&asset.kind);
        ui.painter().text(
            thumb_rect.center(),
            egui::Align2::CENTER_CENTER,
            icon.as_str(),
            egui::FontId::new(28.0, egui::FontFamily::Name("phosphor-regular".into())),
            crate::ui::theme::TEXT,
        );
    }

    ui.painter().text(
        egui::pos2(rect.center().x, rect.max.y - 6.0),
        egui::Align2::CENTER_BOTTOM,
        asset.label.as_str(),
        egui::FontId::new(11.0, egui::FontFamily::Proportional),
        if selected {
            crate::ui::theme::TEXT_SELECTED
        } else {
            crate::ui::theme::TEXT
        },
    );

    response
}

#[derive(Clone)]
pub struct ThumbnailCacheSnapshot {
    pub texture_ids: HashMap<String, egui::TextureId>,
    pub pending_keys: Vec<String>,
    pub failed_keys: Vec<String>,
}

impl ThumbnailCacheSnapshot {
    pub fn texture_for_key(&self, key: &str) -> Option<egui::TextureId> {
        self.texture_ids.get(key).copied()
    }

    pub fn is_pending_key(&self, key: &str) -> bool {
        self.pending_keys.iter().any(|pending| pending == key)
    }

    pub fn is_failed_key(&self, key: &str) -> bool {
        self.failed_keys.iter().any(|failed| failed == key)
    }

    pub fn texture_for(&self, asset: &EditorAsset) -> Option<egui::TextureId> {
        self.texture_ids.get(&asset_cache_key(asset)).copied()
    }

    pub fn is_pending(&self, asset: &EditorAsset) -> bool {
        self.pending_keys
            .iter()
            .any(|key| key == &asset_cache_key(asset))
    }

    pub fn is_failed(&self, asset: &EditorAsset) -> bool {
        self.failed_keys
            .iter()
            .any(|key| key == &asset_cache_key(asset))
    }
}

pub fn prefetch_folder_thumbnails(world: &mut World, folder: &str) {
    {
        let cache = world.resource::<AssetThumbnailCache>();
        if cache.prefetched_folder.as_deref() == Some(folder) {
            return;
        }
    }

    let requests: Vec<(String, String, EditorAssetKind)> = {
        let assets = world.resource::<EditorAssets>();
        assets
            .assets
            .iter()
            .filter(|asset| asset.folder_path == folder)
            .filter(|asset| {
                matches!(
                    asset.kind,
                    EditorAssetKind::Texture | EditorAssetKind::Model | EditorAssetKind::Material
                )
            })
            .filter_map(|asset| {
                Some((
                    asset_cache_key(asset),
                    asset.path.clone()?,
                    asset.kind.clone(),
                ))
            })
            .collect()
    };
    let asset_server = world.resource::<AssetServer>().clone();
    world.resource_scope(|world, mut cache: Mut<AssetThumbnailCache>| {
        world.resource_scope(|_world, mut studio: Mut<ThumbnailStudio>| {
            for (key, path, kind) in requests {
                match kind {
                    EditorAssetKind::Texture => cache.request_texture(key, path, &asset_server),
                    EditorAssetKind::Model => {
                        cache.request_model(key, path, &asset_server, &mut studio)
                    }
                    EditorAssetKind::Material => {
                        cache.request_material_asset(key, path, &mut studio)
                    }
                    _ => {}
                }
            }
        });
    });
    world
        .resource_mut::<AssetThumbnailCache>()
        .prefetched_folder = Some(folder.to_string());
}

pub fn invalidate_on_catalog_refresh(world: &mut World) {
    world.resource_mut::<AssetThumbnailCache>().invalidate_all();
    if !world.contains_resource::<ThumbnailStudio>() {
        return;
    }
    let cleanup = {
        let mut studio = world.resource_mut::<ThumbnailStudio>();
        studio.clear_queue();
        studio.clear_jobs()
    };
    if let Some(root) = cleanup.active_root {
        world.commands().entity(root).despawn_children();
        world.commands().entity(root).despawn();
    }
}