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

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

use async_std::task;
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::{ImageLoader, LoaderError, ResourceLoader}, resource::ResHandle, ResourceData, ResourceState, UntypedResHandle};

/// 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 version(&self) -> usize;
    fn uuid(&self) -> Uuid;
    fn path(&self) -> Option<String>;
    fn is_watched(&self) -> bool;
    fn is_loaded(&self) -> bool;
    fn set_state(&self, new: ResourceState);
    fn clone_untyped(&self) -> UntypedResHandle;
}

#[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),
    
    #[error("The resource was not loaded from a path so cannot be reloaded")]
    NoReloadPath
}

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

/// A struct that stores some things used for watching resources.
pub struct ResourceWatcher {
    debouncer: Arc<RwLock<notify_debouncer_full::Debouncer<RecommendedWatcher, FileIdMap>>>,
    events_recv: Receiver<Result<Vec<DebouncedEvent>, Vec<notify::Error>>>,
}

/// The state of the ResourceManager
pub struct ResourceManagerState {
    resources: HashMap<String, Arc<dyn ResourceStorage>>,
    uuid_resources: HashMap<Uuid, Arc<dyn ResourceStorage>>,
    loaders: Vec<Arc<dyn ResourceLoader>>,
    watchers: HashMap<String, ResourceWatcher>,
}

/// The ResourceManager
/// 
/// This exists since we need the manager to be `Send + Sync`.
#[derive(Clone)]
pub struct ResourceManager {
    inner: Arc<RwLock<ResourceManagerState>>,
}

impl Default for ResourceManager {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(
                ResourceManagerState {
                    resources: HashMap::new(),
                    uuid_resources: HashMap::new(),
                    loaders: vec![ Arc::new(ImageLoader), Arc::new(ModelLoader) ],
                    watchers: HashMap::new(),
                }
            ))
        }
    }
}

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

    /// Retrieves a non-mutable guard of the manager's state.
    pub fn state(&self) -> RwLockReadGuard<ResourceManagerState> {
        self.inner.read().unwrap()
    }

    /// Retrieves a mutable guard of the manager's state.
    pub fn state_mut(&self) -> RwLockWriteGuard<ResourceManagerState> {
        self.inner.write().unwrap()
    }

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

    /// Request a resource at `path`.
    /// 
    /// When a resource for a path is requested for the first time, the resource will be loaded
    /// and cached. In the future, the cached version will be returned. There is only ever one copy
    /// of the resource's data in memory at a time.
    /// 
    /// Loading resources is done asynchronously on a task spawned by `async-std`. You can use the
    /// handle to check if the resource is loaded.
    #[inline(always)]
    pub fn request<T>(&self, path: &str) -> Result<ResHandle<T>, RequestError>
    where
        T: ResourceData
    {
        self.request_raw(path)
            .map(|res| res.as_typed::<T>()
                .expect("mismatched asset type, cannot downcast"))
    }

    /// Request a resource without downcasting to a `ResHandle<T>`.
    /// Whenever you're ready to downcast, you can do so like this:
    /// ```nobuild
    /// 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(&self, path: &str) -> Result<UntypedResHandle, RequestError> {
        let mut state = self.state_mut();
        match state.resources.get(&path.to_string()) {
            Some(res) => {
                Ok(res.clone().clone_untyped())
            },
            None => {
                if let Some(loader) = state.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.clone(), path);

                    let handle = loader.create_erased_handle();

                    let untyped = handle.clone_untyped();
                    untyped.write().path = Some(path.to_string());

                    task::spawn(async move {
                        match res.await {
                            Ok(data) => {
                                let mut d = untyped.write();
                                d.state = ResourceState::Ready(data);
                                d.condvar.1.notify_all();
                            }
                            Err(err) => {
                                let mut d = untyped.write();
                                d.state = ResourceState::Error(Arc::new(err));
                            }
                        }
                    });

                    let res: Arc<dyn ResourceStorage> = Arc::from(handle.clone());
                    state.resources.insert(path.to_string(), res.clone());
                    state.uuid_resources.insert(res.uuid(), res);

                    Ok(handle.clone_untyped())
                } 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: ResourceData>(&self, res: ResHandle<T>) {
        let mut state = self.state_mut();
        state.resources.insert(res.uuid().to_string(), 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: ResourceData>(&self, uuid: &Uuid) -> Option<ResHandle<T>> {
        let state = self.state();
        match state.resources.get(&uuid.to_string())
            .or_else(|| state.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>(&self, ident: &str, mime_type: &str, bytes: Vec<u8>, offset: usize, length: usize) -> Result<ResHandle<T>, RequestError>
    where
        T: ResourceData
    {
        let mut state = self.state_mut();
        if let Some(loader) = state.loaders.iter()
                .find(|l| l.does_support_mime(mime_type)) {
            let loader = loader.clone();
            let res = loader.load_bytes(self.clone(), bytes, offset, length);

            let handle = ResHandle::<T>::new_loading(None);
            let thand = handle.clone();
            task::spawn(async move {
                match res.await {
                    Ok(data) => {
                        let mut d = thand.write();
                        d.state = ResourceState::Ready(data);
                    }
                    Err(err) => {
                        let mut d = thand.write();
                        d.state = ResourceState::Error(Arc::new(err));
                    }
                }
            });

            let res: Arc<dyn ResourceStorage> = Arc::from(handle.clone());
            state.resources.insert(ident.to_string(), res.clone());
            state.uuid_resources.insert(res.uuid(), res);

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

    /// Requests bytes from the manager.
    pub fn request_loaded_bytes<T>(&self, ident: &str) -> Result<Arc<ResHandle<T>>, RequestError>
    where
        T: ResourceData
    {
        let state = self.state();
        match state.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(&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(),
        };

        let mut state = self.state_mut();
        state.watchers.insert(path.to_string(), watcher);

        let res = state.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(&self, path: &str) -> notify::Result<()> {
        let state = self.state();
        if let Some(watcher) = state.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 = state.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>>>> {
        let state = self.state();
        state.watchers.get(&path.to_string())
            .map(|w| w.events_recv.clone())
    }

    /// Trigger a reload of a resource.
    /// 
    /// The version of the resource will be incremented by one.
    /// 
    /// > Note: Since reloading is done asynchronously, the reloaded data will not be immediately
    /// accessible. Until the resource is reloaded, the previous data will stay inside of
    /// the handle.
    pub fn reload<T>(&self, resource: ResHandle<T>) -> Result<(), RequestError>
    where
        T: ResourceData
    {
        let state = self.state();

        let path = resource.path()
            .ok_or(RequestError::NoReloadPath)?;
        if let Some(loader) = state.loaders.iter()
                .find(|l| l.does_support_file(&path)) {
            let loader = Arc::clone(loader); // stop borrowing from self
            let res = loader.load(self.clone(), &path);

            /* let res_lock = &resource.data;
            let mut res_lock = res_lock.write().unwrap();
            res_lock.state = ResourceState::Loading;
            drop(res_lock); */
            
            let thand = resource.clone();
            task::spawn(async move {
                match res.await {
                    Ok(data) => {
                        let mut d = thand.write();
                        d.state = ResourceState::Ready(data);
                        d.version += 1;
                    }
                    Err(err) => {
                        let mut d = thand.write();
                        d.state = ResourceState::Error(Arc::new(err));
                    }
                }
            });
        }

        Ok(())
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use std::{io, ops::Deref};

    use instant::Instant;

    use crate::{Image, ResourceData};

    use super::*;

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

    pub(crate) fn busy_wait_resource<R: ResourceData>(handle: &ResHandle<R>, timeout: f32) {
        let start = Instant::now();
        while !handle.is_loaded() {
            // loop until the image is loaded
            let now = Instant::now();

            let diff = now - start;

            if diff.as_secs_f32() >= timeout {
                panic!("Image never loaded");
            }
        }
    }

    pub(crate) fn busy_wait_resource_reload<R: ResourceData>(handle: &ResHandle<R>, timeout: f32) {
        let version = handle.version();
        let start = Instant::now();

        while !handle.is_loaded() || handle.version() == version {
            // loop until the image is loaded
            let now = Instant::now();

            let diff = now - start;

            if diff.as_secs_f32() >= timeout {
                panic!("Image never loaded");
            }
        }
    }

    #[test]
    fn load_image() {
        let man = ResourceManager::new();
        let res = man.request::<Image>(&get_image("squiggles.png")).unwrap();
        assert!(!res.is_loaded());

        res.wait_for_load();
        //busy_wait_resource(&res, 10.0);

        // shouldn't panic because of the loop
        res.data_ref().unwrap();
    }

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

        let resagain = man.request::<Image>(&get_image("squiggles.png")).unwrap();
        assert_eq!(Arc::strong_count(&resagain.handle.res), 4);
    }

    /// Ensures that an error is returned when a file that doesn't exist is requested
    #[test]
    fn ensure_none() {
        let man = ResourceManager::new();
        let res = man.request::<Image>(&get_image("squigglesfff.png")).unwrap();
        //let err = res.err().unwrap();

        // 1 second should be enough to run into an error
        std::thread::sleep(Duration::from_secs(1));
        //busy_wait_resource(&res, 10.0);
        let state = &res.read().state;
        
        assert!(
            match state {
                // make sure the error is NotFound
                //RequestError::Loader(LoaderError::IoError(e)) if e.kind() == io::ErrorKind::NotFound => true,
                ResourceState::Error(err) => {
                    match err.deref() {
                        LoaderError::IoError(e) if e.kind() == io::ErrorKind::NotFound => true,
                        _ => false,
                    }
                },
                _ => false
            }
        );
    }

    #[test]
    fn reload_image() {
        let man = ResourceManager::new();
        let res = man.request::<Image>(&get_image("squiggles.png")).unwrap();
        busy_wait_resource(&res, 10.0);
        let img = res.data_ref();
        img.unwrap();

        man.reload(res.clone()).unwrap();
        busy_wait_resource_reload(&res, 10.0);
        assert_eq!(res.version(), 1);

        man.reload(res.clone()).unwrap();
        busy_wait_resource_reload(&res, 10.0);
        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 man = ResourceManager::new();
        let res = man.request::<Image>(&image_path).unwrap();
        busy_wait_resource(&res, 10.0);
        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()));
    }
}