path: root/engine/src/asset.rs

use std::any::{type_name, Any};
use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::HashMap;
use std::convert::Infallible;
use std::ffi::{OsStr, OsString};
use std::fmt::{Debug, Display};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::sync::mpsc::{
    channel as mpsc_channel,
    Receiver as MpscReceiver,
    Sender as MpscSender,
};
use std::sync::Arc;

use ecs::phase::PRE_UPDATE as PRE_UPDATE_PHASE;
use ecs::sole::Single;
use ecs::Sole;

use crate::work_queue::{Work, WorkQueue};

/// Asset label.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Label<'a>
{
    pub path: Cow<'a, Path>,
    pub name: Option<Cow<'a, str>>,
}

impl Label<'_>
{
    pub fn to_owned(&self) -> LabelOwned
    {
        LabelOwned {
            path: self.path.to_path_buf(),
            name: self.name.as_ref().map(|name| name.to_string()),
        }
    }
}

impl<'a> From<&'a Path> for Label<'a>
{
    fn from(path: &'a Path) -> Self
    {
        Self { path: path.into(), name: None }
    }
}

impl From<PathBuf> for Label<'_>
{
    fn from(path: PathBuf) -> Self
    {
        Self { path: path.into(), name: None }
    }
}

impl<'a> From<&'a LabelOwned> for Label<'a>
{
    fn from(label: &'a LabelOwned) -> Self
    {
        Self {
            path: (&label.path).into(),
            name: label.name.as_ref().map(|name| Cow::Borrowed(name.as_str())),
        }
    }
}

impl Display for Label<'_>
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        write!(formatter, "{}", self.path.display())?;

        if let Some(name) = &self.name {
            formatter.write_str("::")?;
            formatter.write_str(&name)?;
        }

        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct LabelOwned
{
    pub path: PathBuf,
    pub name: Option<String>,
}

impl LabelOwned
{
    pub fn to_label(&self) -> Label<'_>
    {
        Label {
            path: (&self.path).into(),
            name: self.name.as_ref().map(|name| Cow::Borrowed(name.as_str())),
        }
    }
}

impl Display for LabelOwned
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        write!(formatter, "{}", self.path.display())?;

        if let Some(name) = &self.name {
            formatter.write_str("::")?;
            formatter.write_str(&name)?;
        }

        Ok(())
    }
}

#[derive(Debug, Sole)]
pub struct Assets
{
    assets: Vec<StoredAsset>,
    asset_lookup: RefCell<HashMap<LabelHash, LookupEntry>>,
    importers: Vec<WrappedImporterFn>,
    importer_lookup: HashMap<OsString, usize>,
    import_work_queue: WorkQueue<ImportWorkUserData>,
    import_work_msg_receiver: MpscReceiver<ImportWorkMessage>,
    import_work_msg_sender: MpscSender<ImportWorkMessage>,
}

impl Assets
{
    pub fn with_capacity(capacity: usize) -> Self
    {
        let (import_work_msg_sender, import_work_msg_receiver) =
            mpsc_channel::<ImportWorkMessage>();

        Self {
            assets: Vec::with_capacity(capacity),
            asset_lookup: RefCell::new(HashMap::with_capacity(capacity)),
            importers: Vec::new(),
            importer_lookup: HashMap::new(),
            import_work_queue: WorkQueue::new(),
            import_work_msg_receiver,
            import_work_msg_sender,
        }
    }

    pub fn set_importer<'file_ext, AssetSettings, Err>(
        &mut self,
        file_extensions: impl IntoIterator<Item: Into<Cow<'file_ext, str>>>,
        func: impl Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>,
    ) where
        AssetSettings: 'static,
        Err: std::error::Error + 'static,
    {
        self.importers.push(WrappedImporterFn::new(func));

        let importer_index = self.importers.len() - 1;

        self.importer_lookup
            .extend(file_extensions.into_iter().map(|file_ext| {
                let file_ext: Cow<str> = file_ext.into();

                (file_ext.into_owned().into(), importer_index)
            }));
    }

    #[tracing::instrument(skip_all, fields(asset_type=type_name::<Asset>()))]
    pub fn get<'this, 'handle, Asset: 'static + Send + Sync>(
        &'this self,
        handle: &'handle Handle<Asset>,
    ) -> Option<&'handle Asset>
    where
        'this: 'handle,
    {
        let LookupEntry::Occupied(asset_index) =
            *self.asset_lookup.borrow().get(&handle.id.label_hash)?
        else {
            return None;
        };

        let stored_asset = self.assets.get(asset_index).expect("Not possible");

        let Some(asset) = stored_asset.strong.downcast_ref::<Asset>() else {
            tracing::error!("Wrong asset type");
            return None;
        };

        Some(asset)
    }

    #[tracing::instrument(skip(self))]
    pub fn load<'i, Asset: 'static + Send + Sync>(
        &self,
        label: impl Into<Label<'i>> + Debug,
    ) -> Handle<Asset>
    {
        let label = label.into();

        let label_hash = LabelHash::new(&label);

        let mut asset_lookup = self.asset_lookup.borrow_mut();

        if Self::is_pending(&asset_lookup, &label) {
            return Handle::new(label_hash);
        }

        let Some(lookup_entry) = asset_lookup.get(&label_hash) else {
            self.add_import_work::<Infallible>(
                &label,
                label_hash,
                None,
                &mut asset_lookup,
            );

            return Handle::new(label_hash);
        };

        match *lookup_entry {
            LookupEntry::Occupied(asset_index) => {
                let stored_asset = self.assets.get(asset_index).expect("Not possible");

                if stored_asset.strong.downcast_ref::<Asset>().is_none() {
                    tracing::error!("Wrong asset type {}", type_name::<Asset>());
                }
            }
            LookupEntry::Pending => {}
        }

        Handle::new(label_hash)
    }

    #[tracing::instrument(skip(self))]
    pub fn load_with_settings<'i, Asset, AssetSettings>(
        &self,
        label: impl Into<Label<'i>> + Debug,
        asset_settings: AssetSettings,
    ) -> Handle<Asset>
    where
        Asset: Send + Sync + 'static,
        AssetSettings: Send + Sync + Debug + 'static,
    {
        let label = label.into();

        let label_hash = LabelHash::new(&label);

        let mut asset_lookup = self.asset_lookup.borrow_mut();

        if Self::is_pending(&asset_lookup, &label) {
            return Handle::new(label_hash);
        }

        let Some(lookup_entry) = asset_lookup.get(&label_hash) else {
            self.add_import_work::<AssetSettings>(
                &label,
                label_hash,
                Some(asset_settings),
                &mut asset_lookup,
            );

            return Handle::new(label_hash);
        };

        match *lookup_entry {
            LookupEntry::Occupied(asset_index) => {
                let stored_asset = self.assets.get(asset_index).expect("Not possible");

                if stored_asset.strong.downcast_ref::<Asset>().is_none() {
                    tracing::error!(
                        "Wrong asset type {} for asset",
                        type_name::<Asset>()
                    );
                }
            }
            LookupEntry::Pending => {}
        }

        Handle::new(label_hash)
    }

    pub fn store_with_name<'name, Asset: 'static + Send + Sync>(
        &mut self,
        name: impl Into<Cow<'name, str>>,
        asset: Asset,
    ) -> Handle<Asset>
    {
        self.store_with_label(
            Label {
                path: Path::new("").into(),
                name: Some(name.into()),
            },
            asset,
        )
    }

    #[tracing::instrument(skip(self, asset), fields(asset_type=type_name::<Asset>()))]
    pub fn store_with_label<'i, Asset: 'static + Send + Sync>(
        &mut self,
        label: impl Into<Label<'i>> + Debug,
        asset: Asset,
    ) -> Handle<Asset>
    {
        let label = label.into();

        let label_hash = LabelHash::new(&label);

        if matches!(
            self.asset_lookup.get_mut().get(&label_hash),
            Some(LookupEntry::Occupied(_))
        ) {
            tracing::error!("Asset already exists");

            return Handle::new(label_hash);
        }

        tracing::debug!("Storing asset");

        self.assets.push(StoredAsset::new(asset));

        let index = self.assets.len() - 1;

        self.asset_lookup
            .get_mut()
            .insert(label_hash, LookupEntry::Occupied(index));

        if label.name.is_some() {
            let parent_asset_label_hash =
                LabelHash::new(&Label { path: label.path, name: None });

            if matches!(
                self.asset_lookup.get_mut().get(&parent_asset_label_hash),
                Some(LookupEntry::Pending)
            ) {
                self.asset_lookup.get_mut().remove(&parent_asset_label_hash);
            } else if self
                .asset_lookup
                .get_mut()
                .get(&parent_asset_label_hash)
                .is_none()
            {
                self.assets
                    .push(StoredAsset::new::<Option<Infallible>>(None));

                self.asset_lookup.get_mut().insert(
                    parent_asset_label_hash,
                    LookupEntry::Occupied(self.assets.len() - 1),
                );
            }
        }

        Handle::new(label_hash)
    }

    fn is_pending(asset_lookup: &HashMap<LabelHash, LookupEntry>, label: &Label) -> bool
    {
        if label.name.is_some() {
            if let Some(LookupEntry::Pending) =
                asset_lookup.get(&LabelHash::new(&Label {
                    path: label.path.as_ref().into(),
                    name: None,
                }))
            {
                return true;
            }
        }

        if let Some(LookupEntry::Pending) = asset_lookup.get(&LabelHash::new(label)) {
            return true;
        };

        false
    }

    fn add_import_work<AssetSettings>(
        &self,
        label: &Label<'_>,
        label_hash: LabelHash,
        asset_settings: Option<AssetSettings>,
        asset_lookup: &mut HashMap<LabelHash, LookupEntry>,
    ) where
        AssetSettings: Any + Send + Sync,
    {
        let Some(file_ext) = label.path.extension() else {
            tracing::error!("Asset file is missing a file extension");
            return;
        };

        let Some(importer) = self.get_importer(file_ext) else {
            tracing::error!(
                "No importer exists for asset file extension {}",
                file_ext.to_string_lossy()
            );
            return;
        };

        self.import_work_queue.add_work(Work {
            func: |ImportWorkUserData {
                       import_work_msg_sender,
                       asset_path,
                       asset_settings,
                       importer,
                   }| {
                if let Err(err) = importer.call(
                    import_work_msg_sender,
                    asset_path.as_path(),
                    asset_settings.as_deref(),
                ) {
                    tracing::error!(
                        "Failed to load asset {}: {err}",
                        asset_path.display()
                    );
                }
            },
            user_data: ImportWorkUserData {
                import_work_msg_sender: self.import_work_msg_sender.clone(),
                asset_path: label.path.to_path_buf(),
                asset_settings: asset_settings.map(|asset_settings| {
                    Box::new(asset_settings) as Box<dyn Any + Send + Sync>
                }),
                importer: importer.clone(),
            },
        });

        asset_lookup.insert(label_hash, LookupEntry::Pending);

        if label.name.is_some() {
            asset_lookup.insert(
                LabelHash::new(&Label {
                    path: label.path.as_ref().into(),
                    name: None,
                }),
                LookupEntry::Pending,
            );
        }
    }

    fn get_importer(&self, file_ext: &OsStr) -> Option<&WrappedImporterFn>
    {
        let index = *self.importer_lookup.get(file_ext)?;

        Some(self.importers.get(index).expect("Not possible"))
    }
}

impl Default for Assets
{
    fn default() -> Self
    {
        Self::with_capacity(0)
    }
}

pub struct Submitter<'path>
{
    import_work_msg_sender: MpscSender<ImportWorkMessage>,
    asset_path: &'path Path,
}

impl Submitter<'_>
{
    pub fn submit_load_other<'label, Asset: Send + Sync + 'static>(
        &self,
        label: impl Into<Label<'label>>,
    ) -> Handle<Asset>
    {
        let label = label.into();

        let _ = self.import_work_msg_sender.send(ImportWorkMessage::Load {
            do_load: |assets, label, _asset_settings| {
                let _ = assets.load::<Asset>(label);
            },
            label: label.to_owned(),
            asset_settings: None,
        });

        Handle::new(LabelHash::new(&label))
    }

    pub fn submit_load_other_with_settings<'label, Asset, AssetSettings>(
        &self,
        label: impl Into<Label<'label>>,
        asset_settings: AssetSettings,
    ) -> Handle<Asset>
    where
        Asset: Send + Sync + 'static,
        AssetSettings: Send + Sync + Debug + 'static,
    {
        let label = label.into();

        let _ = self.import_work_msg_sender.send(ImportWorkMessage::Load {
            do_load: |assets, label, asset_settings| {
                let asset_settings = *asset_settings
                    .expect("Not possible")
                    .downcast::<AssetSettings>()
                    .expect("Not possible");

                let _ = assets
                    .load_with_settings::<Asset, AssetSettings>(label, asset_settings);
            },
            label: label.to_owned(),
            asset_settings: Some(Box::new(asset_settings)),
        });

        Handle::new(LabelHash::new(&label))
    }

    pub fn submit_store<Asset: Send + Sync + 'static>(
        &self,
        asset: Asset,
    ) -> Handle<Asset>
    {
        let label = LabelOwned {
            path: self.asset_path.into(),
            name: None,
        };

        let label_hash = LabelHash::new(&label.to_label());

        let _ = self.import_work_msg_sender.send(ImportWorkMessage::Store {
            do_store: |assets, label, boxed_asset| {
                let Ok(asset) = boxed_asset.downcast::<Asset>() else {
                    unreachable!();
                };

                assets.store_with_label::<Asset>(&label, *asset);
            },
            label,
            asset: Box::new(asset),
        });

        Handle::new(label_hash)
    }

    pub fn submit_store_named<Asset: Send + Sync + 'static>(
        &self,
        name: impl AsRef<str>,
        asset: Asset,
    ) -> Handle<Asset>
    {
        let label = LabelOwned {
            path: self.asset_path.into(),
            name: Some(name.as_ref().into()),
        };

        let label_hash = LabelHash::new(&label.to_label());

        let _ = self.import_work_msg_sender.send(ImportWorkMessage::Store {
            do_store: |assets, label, boxed_asset| {
                let Ok(asset) = boxed_asset.downcast::<Asset>() else {
                    unreachable!();
                };

                assets.store_with_label::<Asset>(&label, *asset);
            },
            label,
            asset: Box::new(asset),
        });

        Handle::new(label_hash)
    }
}

/// Asset handle.
#[derive(Debug)]
pub struct Handle<Asset: 'static>
{
    id: Id,
    _pd: PhantomData<Asset>,
}

impl<Asset: 'static> Handle<Asset>
{
    pub fn id(&self) -> Id
    {
        self.id
    }

    fn new(label_hash: LabelHash) -> Self
    {
        Self {
            id: Id { label_hash },
            _pd: PhantomData,
        }
    }
}

impl<Asset: 'static> Clone for Handle<Asset>
{
    fn clone(&self) -> Self
    {
        Self { id: self.id, _pd: PhantomData }
    }
}

/// Asset ID.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Id
{
    label_hash: LabelHash,
}

#[derive(Debug, thiserror::Error)]
enum ImporterError
{
    #[error("Settings has a incorrect type")]
    IncorrectAssetSettingsType(PathBuf),

    #[error(transparent)]
    Other(Box<dyn std::error::Error>),
}

#[derive(Debug, Clone)]
struct WrappedImporterFn
{
    wrapper_func: fn(
        MpscSender<ImportWorkMessage>,
        &Path,
        Option<&(dyn Any + Send + Sync)>,
    ) -> Result<(), ImporterError>,
}

impl WrappedImporterFn
{
    fn new<InnerFunc, AssetSettings, Err>(inner_func_param: InnerFunc) -> Self
    where
        InnerFunc:
            Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>,
        AssetSettings: 'static,
        Err: std::error::Error + 'static,
    {
        assert_eq!(size_of::<InnerFunc>(), 0);

        let wrapper_func =
            |import_work_msg_sender: MpscSender<ImportWorkMessage>,
             asset_path: &Path,
             asset_settings: Option<&(dyn Any + Send + Sync)>| {
                let inner_func = unsafe { std::mem::zeroed::<InnerFunc>() };

                let asset_settings = asset_settings
                    .map(|asset_settings| {
                        asset_settings
                            .downcast_ref::<AssetSettings>()
                            .ok_or_else(|| {
                                ImporterError::IncorrectAssetSettingsType(
                                    asset_path.to_path_buf(),
                                )
                            })
                    })
                    .transpose()?;

                inner_func(
                    &mut Submitter { import_work_msg_sender, asset_path },
                    asset_path,
                    asset_settings,
                )
                .map_err(|err| ImporterError::Other(Box::new(err)))?;

                Ok(())
            };

        std::mem::forget(inner_func_param);

        Self { wrapper_func }
    }

    fn call(
        &self,
        import_work_msg_sender: MpscSender<ImportWorkMessage>,
        asset_path: &Path,
        asset_settings: Option<&(dyn Any + Send + Sync)>,
    ) -> Result<(), ImporterError>
    {
        (self.wrapper_func)(import_work_msg_sender, asset_path, asset_settings)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct LabelHash(u64);

impl LabelHash
{
    fn new(label: &Label<'_>) -> Self
    {
        let mut hasher = DefaultHasher::new();

        label.hash(&mut hasher);

        Self(hasher.finish())
    }
}

#[derive(Debug, Default)]
pub(crate) struct Extension
{
    pub assets: Assets,
}

impl ecs::extension::Extension for Extension
{
    fn collect(self, mut collector: ecs::extension::Collector<'_>)
    {
        let _ = collector.add_sole(self.assets);

        collector.add_system(*PRE_UPDATE_PHASE, add_received_assets);
    }
}

fn add_received_assets(mut assets: Single<Assets>)
{
    while let Some(import_work_msg) = assets.import_work_msg_receiver.try_recv().ok() {
        match import_work_msg {
            ImportWorkMessage::Store { do_store, label, asset } => {
                do_store(&mut assets, label, asset);
            }
            ImportWorkMessage::Load { do_load, label, asset_settings } => {
                do_load(
                    &assets,
                    Label {
                        path: label.path.as_path().into(),
                        name: label.name.as_deref().map(|name| name.into()),
                    },
                    asset_settings,
                );
            }
        }
    }
}

#[derive(Debug)]
struct ImportWorkUserData
{
    import_work_msg_sender: MpscSender<ImportWorkMessage>,
    asset_path: PathBuf,
    asset_settings: Option<Box<dyn Any + Send + Sync>>,
    importer: WrappedImporterFn,
}

#[derive(Debug)]
enum ImportWorkMessage
{
    Store
    {
        do_store: fn(&mut Assets, LabelOwned, Box<dyn Any + Send + Sync>),
        label: LabelOwned,
        asset: Box<dyn Any + Send + Sync>,
    },

    Load
    {
        do_load: fn(&Assets, Label<'_>, Option<Box<dyn Any + Send + Sync>>),
        label: LabelOwned,
        asset_settings: Option<Box<dyn Any + Send + Sync>>,
    },
}

#[derive(Debug, Clone, Copy)]
enum LookupEntry
{
    Occupied(usize),
    Pending,
}

#[derive(Debug)]
struct StoredAsset
{
    strong: Arc<dyn Any + Send + Sync>,
}

impl StoredAsset
{
    fn new<Asset: Any + Send + Sync>(asset: Asset) -> Self
    {
        let strong = Arc::new(asset);

        Self { strong }
    }
}