lyra-engine/lyra-resource/src/resource_manager.rs

use std::{sync::{Arc, RwLock}, collections::HashMap, any::Any, path::Path, time::Duration};

use crossbeam::channel::Receiver;
use notify::{Watcher, RecommendedWatcher};
use notify_debouncer_full::{DebouncedEvent, FileIdMap};
use thiserror::Error;
use uuid::Uuid;

use crate::{gltf::ModelLoader, loader::{image::ImageLoader, LoaderError, ResourceLoader}, resource::ResHandle, ResourceState};

/// A trait for type erased storage of a resource.
/// Implemented for [`ResHandle<T>`]
pub trait ResourceStorage: Send + Sync + Any + 'static {
    fn as_any(&self) -> &dyn Any;
    fn as_any_mut(&mut self) -> &mut dyn Any;
    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
    fn as_box_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
    /// Do not set a resource to watched if it is not actually watched.
    /// This is used internally.
    fn set_watched(&self, watched: bool);

    fn path(&self) -> String;
    fn version(&self) -> usize;
    fn state(&self) -> ResourceState;
    fn uuid(&self) -> Uuid;
    fn is_watched(&self) -> bool;
    fn is_loaded(&self) -> bool;
}

#[derive(Error, Debug)]
pub enum RequestError {
    #[error("{0}")]
    Loader(LoaderError),
    #[error("The file extension is unsupported: '{0}'")]
    UnsupportedFileExtension(String),
    #[error("The mimetype is unsupported: '{0}'")]
    UnsupportedMime(String),
    #[error("The identifier is not found: '{0}'")]
    IdentNotFound(String),
}

impl From<LoaderError> for RequestError {
    fn from(value: LoaderError) -> Self {
        RequestError::Loader(value)
    }
}

/// A struct that stores all Manager data. This is requried for sending 
//struct ManagerStorage

/// A struct that
pub struct ResourceWatcher {
    debouncer: Arc<RwLock<notify_debouncer_full::Debouncer<RecommendedWatcher, FileIdMap>>>,
    events_recv: Receiver<Result<Vec<DebouncedEvent>, Vec<notify::Error>>>,
}

pub struct ResourceManager {
    resources: HashMap<String, Arc<dyn ResourceStorage>>,
    uuid_resources: HashMap<Uuid, Arc<dyn ResourceStorage>>,
    loaders: Vec<Arc<dyn ResourceLoader>>,
    watchers: HashMap<String, ResourceWatcher>,
}

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

impl ResourceManager {
    pub fn new() -> Self {
        Self {
            resources: HashMap::new(),
            uuid_resources: HashMap::new(),
            loaders: vec![ Arc::new(ImageLoader), Arc::new(ModelLoader) ],
            watchers: HashMap::new(),
        }
    }

    /// Registers a loader to the manager.
    pub fn register_loader<L>(&mut self)
    where
        L: ResourceLoader + Default + 'static
    {
        self.loaders.push(Arc::new(L::default()));
    }

    pub fn request<T>(&mut self, path: &str) -> Result<ResHandle<T>, RequestError>
    where
        T: Send + Sync + Any + 'static
    {
        match self.resources.get(&path.to_string()) {
            Some(res) => {
                let res = res.clone().as_arc_any();
                let res: Arc<ResHandle<T>> = res.downcast::<ResHandle<T>>().expect("Failure to downcast resource");
                let res = ResHandle::<T>::clone(&res);

                Ok(res)
            },
            None => {
                if let Some(loader) = self.loaders.iter()
                        .find(|l| l.does_support_file(path)) {

                    // Load the resource and store it
                    let loader = Arc::clone(loader); // stop borrowing from self
                    let res = loader.load(self, path)?;
                    let res: Arc<dyn ResourceStorage> = Arc::from(res);
                    self.resources.insert(path.to_string(), res.clone());

                    // cast Arc<dyn ResourceStorage> to Arc<Resource<T>
                    let res = res.as_arc_any();
                    let res = res.downcast::<ResHandle<T>>()
                        .expect("Failure to downcast resource! Does the loader return an `Arc<Resource<T>>`?");
                    let res = ResHandle::<T>::clone(&res);

                    Ok(res)
                } else {
                    Err(RequestError::UnsupportedFileExtension(path.to_string()))
                }
            }
        }
    }

    /// Request a resource without downcasting to a ResHandle<T>.
    /// Whenever you're ready to downcast, you can do so like this:
    /// ```compile_fail
    /// let arc_any = res_arc.as_arc_any();
    /// let res: Arc<ResHandle<T>> = res.downcast::<ResHandle<T>>().expect("Failure to downcast resource");
    /// ```
    pub fn request_raw(&mut self, path: &str) -> Result<Arc<dyn ResourceStorage>, RequestError> {
        match self.resources.get(&path.to_string()) {
            Some(res) => {
                Ok(res.clone())
            },
            None => {
                if let Some(loader) = self.loaders.iter()
                        .find(|l| l.does_support_file(path)) {

                    // Load the resource and store it
                    let loader = Arc::clone(loader); // stop borrowing from self
                    let res = loader.load(self, path)?;
                    let res: Arc<dyn ResourceStorage> = Arc::from(res);
                    self.resources.insert(path.to_string(), res.clone());

                    Ok(res)
                } else {
                    Err(RequestError::UnsupportedFileExtension(path.to_string()))
                }
            }
        }
    }

    /// Store a resource using its uuid.
    /// 
    /// The resource cannot be requested with [`ResourceManager::request`], it can only be 
    /// retrieved with [`ResourceManager::request_uuid`].
    pub fn store_uuid<T: Send + Sync + 'static>(&mut self, res: ResHandle<T>) {
        self.uuid_resources.insert(res.uuid(), Arc::new(res));
    }

    /// Request a resource via its uuid.
    /// 
    /// Returns `None` if the resource was not found. The resource must of had been
    /// stored with [`ResourceManager::request`] to return `Some`.
    pub fn request_uuid<T: Send + Sync + 'static>(&mut self, uuid: &Uuid) -> Option<ResHandle<T>> {
        match self.uuid_resources.get(uuid) {
            Some(res) => {
                let res = res.clone().as_arc_any();
                let res: Arc<ResHandle<T>> = res.downcast::<ResHandle<T>>().expect("Failure to downcast resource");
                Some(ResHandle::<T>::clone(&res))
            },
            None => None,
        }
    }

    /// Store bytes in the manager. If there is already an entry with the same identifier it will be updated.
    /// 
    /// Panics: If there is already an entry with the same `ident`, and the entry is not bytes, this function will panic.
    /// 
    /// Parameters:
    /// * `ident` - The identifier to store along with these bytes. Make sure its unique to avoid overriding something.
    /// * `bytes` - The bytes to store.
    /// 
    /// Returns: The `Arc` to the now stored resource
    pub fn load_bytes<T>(&mut self, ident: &str, mime_type: &str, bytes: Vec<u8>, offset: usize, length: usize) -> Result<ResHandle<T>, RequestError>
    where
        T: Send + Sync + Any + 'static
    {
        if let Some(loader) = self.loaders.iter()
                .find(|l| l.does_support_mime(mime_type)) {
            let loader = loader.clone();
            let res = loader.load_bytes(self, bytes, offset, length)?;
            let res: Arc<dyn ResourceStorage> = Arc::from(res);
            self.resources.insert(ident.to_string(), res.clone());
            // code here...

            // cast Arc<dyn ResourceStorage> to Arc<Resource<T>
            let res = res.as_arc_any();
            let res = res.downcast::<ResHandle<T>>()
                .expect("Failure to downcast resource! Does the loader return an `Arc<Resource<T>>`?");
            let res = ResHandle::<T>::clone(&res);

            Ok(res)
        } else {
            Err(RequestError::UnsupportedMime(mime_type.to_string()))
        }
    }

    /// Requests bytes from the manager.
    pub fn request_loaded_bytes<T>(&mut self, ident: &str) -> Result<Arc<ResHandle<T>>, RequestError>
    where
        T: Send + Sync + Any + 'static
    {
        match self.resources.get(&ident.to_string()) {
            Some(res) => {
                let res = res.clone().as_arc_any();
                let res = res.downcast::<ResHandle<T>>().expect("Failure to downcast resource");

                Ok(res)
            },
            None => {
                Err(RequestError::IdentNotFound(ident.to_string()))
            }
        }
    }

    /// Start watching a path for changes. Returns a mspc channel that will send events
    pub fn watch(&mut self, path: &str, recursive: bool) -> notify::Result<Receiver<Result<Vec<DebouncedEvent>, Vec<notify::Error>>>> {
        let (send, recv) = crossbeam::channel::bounded(15);
        let mut watcher = notify_debouncer_full::new_debouncer(Duration::from_millis(1000), None, send)?;

        let recurse_mode = match recursive {
            true => notify::RecursiveMode::Recursive,
            false => notify::RecursiveMode::NonRecursive,
        };
        watcher.watcher().watch(path.as_ref(), recurse_mode)?;

        let watcher = Arc::new(RwLock::new(watcher));
        let watcher = ResourceWatcher {
            debouncer: watcher,
            events_recv: recv.clone(),
        };

        self.watchers.insert(path.to_string(), watcher);

        let res = self.resources.get(&path.to_string())
            .expect("The path that was watched has not been loaded as a resource yet");
        res.set_watched(true);

        Ok(recv)
    }

    /// Stops watching a path
    pub fn stop_watching(&mut self, path: &str) -> notify::Result<()> {
        if let Some(watcher) = self.watchers.get(path) {
            let mut watcher = watcher.debouncer.write().unwrap();
            watcher.watcher().unwatch(Path::new(path))?;

            // unwrap is safe since only loaded resources can be watched
            let res = self.resources.get(&path.to_string()).unwrap();
            res.set_watched(false);
        }

        Ok(())
    }

    /// Returns a mspc receiver for watcher events of a specific path. The path must already
    /// be watched with [`ResourceManager::watch`] for this to return `Some`.
    pub fn watcher_event_recv(&self, path: &str) -> Option<Receiver<Result<Vec<DebouncedEvent>, Vec<notify::Error>>>> {
        self.watchers.get(&path.to_string())
            .map(|w| w.events_recv.clone())
    }

    /// Reloads a resource. The data inside the resource will be updated, the state may
    pub fn reload<T>(&mut self, resource: ResHandle<T>) -> Result<(), RequestError>
    where
        T: Send + Sync + Any + 'static
    {
        let path = resource.path();
        if let Some(loader) = self.loaders.iter()
                .find(|l| l.does_support_file(&path)) {
            println!("got loader");
            let loader = Arc::clone(loader); // stop borrowing from self
            let loaded = loader.load(self, &path)?;
            let loaded = loaded.as_arc_any();
            
            let loaded = loaded.downcast::<ResHandle<T>>()
                .unwrap();
            let loaded = match Arc::try_unwrap(loaded) {
                Ok(v) => v,
                Err(_) => panic!("Seems impossible that this would happen, the resource was just loaded!"),
            };
            let loaded = loaded.data;
            let loaded = match Arc::try_unwrap(loaded) {
                Ok(v) => v,
                Err(_) => panic!("Seems impossible that this would happen, the resource was just loaded!"),
            };
            let loaded = loaded.into_inner().unwrap();

            let res_lock = &resource.data;
            let mut res_lock = res_lock.write().unwrap();
            let version = res_lock.version;
            
            res_lock.data = loaded.data;
            res_lock.state = loaded.state;
            res_lock.version = version + 1;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use crate::{Texture, ResourceState};

    use super::*;

    fn get_image(path: &str) -> String {
        let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap();
        
        format!("{manifest}/test_files/img/{path}")
    }

    #[test]
    fn load_image() {
        let mut man = ResourceManager::new();
        let res = man.request::<Texture>(&get_image("squiggles.png")).unwrap();
        assert_eq!(res.state(), ResourceState::Ready);
        let img = res.data_ref();
        img.unwrap();
    }

    /// Ensures that only one copy of the same data was made
    #[test]
    fn ensure_single() {
        let mut man = ResourceManager::new();
        let res = man.request::<Texture>(&get_image("squiggles.png")).unwrap();
        assert_eq!(Arc::strong_count(&res.data), 2);

        let resagain = man.request::<Texture>(&get_image("squiggles.png")).unwrap();
        assert_eq!(Arc::strong_count(&resagain.data), 3);
    }

    /// Ensures that an error is returned when a file that doesn't exist is requested
    #[test]
    fn ensure_none() {
        let mut man = ResourceManager::new();
        let res = man.request::<Texture>(&get_image("squigglesfff.png"));
        let err = res.err().unwrap();
        
        assert!(
            match err {
                // make sure the error is NotFound
                RequestError::Loader(LoaderError::IoError(e)) if e.kind() == io::ErrorKind::NotFound => true,
                _ => false
            }
        );
    }

    #[test]
    fn reload_image() {
        let mut man = ResourceManager::new();
        let res = man.request::<Texture>(&get_image("squiggles.png")).unwrap();
        assert_eq!(res.state(), ResourceState::Ready);
        let img = res.data_ref();
        img.unwrap();

        println!("Path = {}", res.path());
        man.reload(res.clone()).unwrap();
        assert_eq!(res.version(), 1);

        man.reload(res.clone()).unwrap();
        assert_eq!(res.version(), 2);
    }

    #[test]
    fn watch_image() {
        let orig_path = get_image("squiggles.png");
        let image_path = get_image("squiggles_test.png");
        std::fs::copy(orig_path, &image_path).unwrap();

        let mut man = ResourceManager::new();
        let res = man.request::<Texture>(&image_path).unwrap();
        assert_eq!(res.state(), ResourceState::Ready);
        let img = res.data_ref();
        img.unwrap();

        let recv = man.watch(&image_path, false).unwrap();

        std::fs::remove_file(&image_path).unwrap();

        let event = recv.recv().unwrap();
        let event = event.unwrap();
        
        assert!(event.iter().any(|ev| ev.kind.is_remove() || ev.kind.is_modify()));
    }
}