tirbofish/dropbear

use crate::graphics::SharedGraphicsContext;
use crate::model::Model;
use crate::texture::{Texture, TextureBuilder};
use crate::utils::ResourceReference;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::marker::PhantomData;
use std::sync::{Arc, LazyLock};

pub static ASSET_REGISTRY: LazyLock<Arc<RwLock<AssetRegistry>>> =
    LazyLock::new(|| Arc::new(RwLock::new(AssetRegistry::new())));

/// A handle with type [`T`] that provides an index to the [AssetRegistry] contents.
#[derive(Hash, Eq, Debug)]
pub struct Handle<T> {
    pub id: u64,
    _phantom: PhantomData<T>,
}

impl<T> PartialEq for Handle<T> {
    fn eq(&self, other: &Self) -> bool {
        if self.is_null() && other.is_null() {
            return false;
        }
        self.id == other.id
    }
}

impl<T> Copy for Handle<T> {}

impl<T> Clone for Handle<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Handle<T> {
    /// Creates a null handle, for when there is no way to uniquely identify a hash (such as a viewport texture).
    ///
    /// A NULL model would be a model with no vertices or any properties, just a husk.
    /// A NULL texture would throw an error, as it is not possible.
    ///
    /// # Safety
    /// You will want to watch out, as adding this onto the asset registry with a type
    /// where there already is a null handle item, it will be overwritten and data
    /// will not be saved. It is the reason why you will want to consider using the [Self::is_null]
    /// function to verify if the storage of the type has gone through correctly.
    pub const NULL: Self = Self {
        id: 0,
        _phantom: PhantomData,
    };

    /// Creates a new handle with the given ID.
    pub fn new(id: u64) -> Self {
        Self {
            id,
            _phantom: Default::default(),
        }
    }

    /// Returns true if the handle is null.
    pub fn is_null(&self) -> bool {
        self.id == 0
    }
}

pub struct AssetRegistry {
    textures: HashMap<u64, Arc<Texture>>,
    texture_labels: HashMap<String, Handle<Texture>>,

    models: HashMap<u64, Arc<Model>>,
    model_labels: HashMap<String, Handle<Model>>,
}

#[repr(C)]
#[derive(Debug, Clone)]
#[dropbear_macro::repr_c_enum]
pub enum AssetKind {
    Texture,
    Model,
}

/// Common
impl AssetRegistry {
    pub fn new() -> Self {
        let mut result = Self {
            textures: Default::default(),
            texture_labels: Default::default(),
            models: Default::default(),
            model_labels: Default::default(),
        };

        result.add_model(Model {
            hash: 0,
            label: "null".to_string(),
            path: Default::default(),
            meshes: vec![],
            materials: vec![],
            skins: vec![],
            animations: vec![],
            nodes: vec![],
            morph_deltas_buffer: None,
        });

        result
    }

    /// A convenient helper function for hashing a byte slice of data.
    pub(crate) fn hash_bytes(data: &[u8]) -> u64 {
        let mut hasher = DefaultHasher::new();
        data.hash(&mut hasher);
        hasher.finish()
    }

    /// Checks if the asset registry contains a handle with the given hash.
    ///
    /// It will check all different types, so it does not point out specifically where.
    pub fn contains_hash(&self, hash: u64) -> bool {
        self.textures.contains_key(&hash) || self.models.contains_key(&hash)
    }

    /// Checks if the asset registry contains a handle with the given string label.
    ///
    /// It will check all different types, so it does not point out specifically where.
    pub fn contains_label(&self, label: &str) -> bool {
        self.texture_labels.contains_key(label) || self.model_labels.contains_key(label)
    }

    /// Flushes away all unused assets and returns the count of flushed models.
    pub fn flush_unused_with_live_ids(&mut self, live_model_ids: &HashSet<u64>) -> usize {
        log::debug!("Flushing unused assets");
        let mut live_texture_ids: HashSet<u64> = HashSet::new();
        let mut counter = 0;

        // models
        self.models.retain(|id, arc| {
            let is_null = *id == 0;
            let is_protected = arc.label.eq_ignore_ascii_case("light cube");
            let is_live = live_model_ids.contains(id);
            let has_external_arc = Arc::get_mut(arc).is_none();

            if is_null || is_protected || is_live || has_external_arc {
                for mat in &arc.materials {
                    live_texture_ids.insert(mat.diffuse_texture.id);
                    if let Some(h) = mat.normal_texture {
                        live_texture_ids.insert(h.id);
                    }
                    if let Some(h) = mat.emissive_texture {
                        live_texture_ids.insert(h.id);
                    }
                    if let Some(h) = mat.metallic_roughness_texture {
                        live_texture_ids.insert(h.id);
                    }
                    if let Some(h) = mat.occlusion_texture {
                        live_texture_ids.insert(h.id);
                    }
                }
                return true;
            }

            counter += 1;
            false
        });
        self.model_labels
            .retain(|_, handle| self.models.contains_key(&handle.id));

        self.textures.retain(|id, arc| {
            let is_null = *id == 0;
            let is_live = live_texture_ids.contains(id);
            let has_external_arc = Arc::get_mut(arc).is_none();

            if is_null || is_live || has_external_arc {
                return true;
            }

            counter += 1;
            false
        });
        self.texture_labels
            .retain(|_, handle| self.textures.contains_key(&handle.id));

        counter
    }

    /// Flushes away all unused assets using Arc strong-count only (no ECS awareness).
    /// Prefer `flush_unused_with_live_ids` when the ECS world is accessible.
    pub fn flush_unused(&mut self) -> usize {
        self.flush_unused_with_live_ids(&HashSet::new())
    }

    pub fn flush_everything(&mut self) {
        log::debug!("Flushing everything");
        self.models.clear();
        self.model_labels.clear();
        self.textures.clear();
        self.texture_labels.clear();
    }
}

/// Texture stuff
impl AssetRegistry {
    /// Adds a texture and returns a handle.
    ///
    /// This assumes a [Texture] has already been created by you. To create a new texture,
    /// you can use [`Texture::from_bytes`].
    pub fn add_texture(&mut self, texture: Texture) -> Handle<Texture> {
        let handle = texture
            .hash
            .map(|v| Handle::new(v))
            .unwrap_or_else(|| Handle::NULL);
        self.textures.entry(handle.id).or_insert(Arc::new(texture));
        handle
    }

    /// Adds a texture with a label. If the texture already exists (by hash),
    /// returns the existing handle and updates the label to point at it.
    pub fn add_texture_with_label(
        &mut self,
        label: impl Into<String>,
        texture: Texture,
    ) -> Handle<Texture> {
        let handle = self.add_texture(texture);
        self.texture_labels.insert(label.into(), handle.clone());
        handle
    }

    pub fn update_texture(
        &mut self,
        handle: Handle<Texture>,
        new_texture: Texture,
    ) -> Option<Arc<Texture>> {
        self.textures.insert(handle.id, Arc::new(new_texture))
    }

    /// Maps a label to an existing texture handle.
    pub fn label_texture(&mut self, label: impl Into<String>, handle: Handle<Texture>) {
        self.texture_labels.insert(label.into(), handle.clone());
    }

    /// Removes a label from the texture registry, but keeps it in the registry.
    ///
    /// When the label is removed, the [Handle] is still valid.
    pub fn remove_label_texture(&mut self, label: &str) {
        self.texture_labels.remove(label);
    }

    pub fn get_texture(&self, handle: Handle<Texture>) -> Option<Arc<Texture>> {
        self.textures.get(&handle.id).cloned()
    }

    pub fn get_texture_by_label(&self, label: &str) -> Option<Arc<Texture>> {
        self.texture_labels
            .get(label)
            .and_then(|handle| self.textures.get(&handle.id))
            .cloned()
    }

    pub fn get_texture_handle_from_label(&self, label: &str) -> Option<Handle<Texture>> {
        self.texture_labels.get(label).cloned()
    }

    pub fn texture_handle_by_hash(&self, hash: u64) -> Option<Handle<Texture>> {
        self.textures.contains_key(&hash).then(|| Handle::new(hash))
    }

    pub fn grey_texture(&mut self, graphics: Arc<SharedGraphicsContext>) -> Handle<Texture> {
        self.solid_texture_rgba8(
            graphics,
            [128, 128, 128, 255],
            Some(Texture::TEXTURE_FORMAT_BASE.add_srgb_suffix()),
        )
    }

    pub fn solid_texture_rgba8(
        &mut self,
        graphics: Arc<SharedGraphicsContext>,
        rgba: [u8; 4],
        format: Option<wgpu::TextureFormat>,
    ) -> Handle<Texture> {
        let format = format.unwrap_or(Texture::TEXTURE_FORMAT);

        let format_tag = format!("{:?}", format);
        let label = format!(
            "Solid texture [{}, {}, {}, {}] {}",
            rgba[0], rgba[1], rgba[2], rgba[3], format_tag,
        );
        if let Some(handle) = self.get_texture_handle_from_label(label.as_str()) {
            return handle;
        }

        let texture = TextureBuilder::new(&graphics.device)
            .with_raw_pixels(graphics.clone(), &rgba)
            .size(1, 1)
            .format(format)
            .label(label.as_str())
            .build();

        self.add_texture_with_label(label, texture)
    }

    pub fn get_label_from_texture_handle(&self, handle: Handle<Texture>) -> Option<String> {
        self.texture_labels.iter().find_map(|(label, h)| {
            if *h == handle {
                Some(label.clone())
            } else {
                None
            }
        })
    }
}

/// Model stuff
impl AssetRegistry {
    pub fn add_model(&mut self, model: Model) -> Handle<Model> {
        let handle = Handle::new(model.hash);
        self.models.entry(handle.id).or_insert(Arc::new(model));
        handle
    }

    pub fn add_model_with_label(
        &mut self,
        label: impl Into<String>,
        model: Model,
    ) -> Handle<Model> {
        let handle = self.add_model(model);
        self.model_labels.insert(label.into(), handle.clone());
        handle
    }

    pub fn label_model(&mut self, label: impl Into<String>, handle: Handle<Model>) {
        self.model_labels.insert(label.into(), handle.clone());
    }

    pub fn update_model(&mut self, handle: Handle<Model>, model: Model) -> Option<Arc<Model>> {
        if let Some(existing) = self.models.get(&handle.id) {
            if existing.label.eq_ignore_ascii_case("light cube") {
                log::warn!("Attempted to update protected model '{}'", existing.label);
                return None;
            }
        }

        self.models.insert(handle.id, Arc::new(model))
    }

    pub fn get_model(&self, handle: Handle<Model>) -> Option<Arc<Model>> {
        self.models.get(&handle.id).cloned()
    }

    pub fn get_model_by_label(&self, label: &str) -> Option<Arc<Model>> {
        self.model_labels
            .get(label)
            .and_then(|handle| self.models.get(&handle.id))
            .cloned()
    }

    pub fn get_model_handle_from_label(&self, label: &str) -> Option<Handle<Model>> {
        self.model_labels.get(label).cloned()
    }

    pub fn list_models(&self) -> Vec<(Handle<Model>, String, ResourceReference)> {
        self.models
            .values()
            .map(|model| {
                (
                    Handle::new(model.hash),
                    model.label.clone(),
                    model.path.clone(),
                )
            })
            .collect()
    }

    pub fn get_model_handle_by_reference(
        &self,
        reference: &ResourceReference,
    ) -> Option<Handle<Model>> {
        self.models
            .values()
            .find(|model| &model.path == reference)
            .map(|model| Handle::new(model.hash))
    }

    pub fn model_handle_by_hash(&self, hash: u64) -> Option<Handle<Model>> {
        self.models.contains_key(&hash).then(|| Handle::new(hash))
    }

    pub fn get_label_from_model_handle(&self, handle: Handle<Model>) -> Option<String> {
        self.model_labels.iter().find_map(|(label, h)| {
            if *h == handle {
                Some(label.clone())
            } else {
                None
            }
        })
    }

    /// Returns a list of all loaded textures with their handles, labels, and references.
    pub fn list_textures(
        &self,
    ) -> Vec<(Handle<Texture>, Option<String>, Option<ResourceReference>)> {
        self.textures
            .values()
            .map(|texture| {
                (
                    texture.hash.map(Handle::new).unwrap_or(Handle::NULL),
                    texture.label.clone(),
                    texture.reference.clone(),
                )
            })
            .collect()
    }

    /// Finds a texture handle by its resource reference.
    pub fn get_texture_handle_by_reference(
        &self,
        reference: &ResourceReference,
    ) -> Option<Handle<Texture>> {
        self.textures
            .values()
            .find(|texture| texture.reference.as_ref() == Some(reference))
            .and_then(|texture| texture.hash.map(Handle::new))
    }
}

impl Default for AssetRegistry {
    fn default() -> Self {
        Self::new()
    }
}