roast2d_internal 0.4.0

use crossbeam_channel::{Receiver, Sender, unbounded};
use hashbrown::HashMap;
use image::DynamicImage;
use std::{
    mem,
    path::{Path, PathBuf},
    sync::Arc,
};

use anyhow::{Result, anyhow};
use io::{ErasedAssetReader, get_default_reader};

use crate::{
    font::Font,
    handle::{DropEvent, Handle, HandleId},
};

/// Warning threshold for asset count to detect potential memory leaks
const ASSET_COUNT_WARNING_THRESHOLD: usize = 4096;

mod io;

#[derive(Debug)]
pub enum AssetType {
    Raw,
    Texture,
    Font,
}

#[derive(Debug)]
pub(crate) struct PendingTask {
    pub handle: Handle,
    pub asset_type: AssetType,
    pub path: PathBuf,
}

#[derive(Debug)]
pub(crate) struct AssetFetchTask {
    pub handle: Handle,
    pub asset_type: AssetType,
    pub path: String,
}

pub(crate) struct AssetFetchRequest {
    pub reader: Arc<dyn ErasedAssetReader>,
    pub tasks: Vec<AssetFetchTask>,
}

#[derive(Debug)]
pub(crate) struct AssetFetchOutput {
    pub handle: Handle,
    pub asset_type: AssetType,
    pub bytes: Vec<u8>,
}

pub(crate) enum FetchedTask {
    CreateTexture { handle: Handle, data: DynamicImage },
    RemoveTexture { handle: u64 },
    CreateFont { handle: Handle, font: Font },
    RemoveFont { handle: u64 },
}

pub struct Asset {
    pub asset_type: AssetType,
    pub bytes: Option<Vec<u8>>,
}

pub struct AssetManager {
    reader: Arc<dyn ErasedAssetReader>,
    pending: Vec<PendingTask>,
    pending_fetched: Vec<FetchedTask>,
    asset_id: u64,
    receiver: Receiver<DropEvent>,
    sender: Sender<DropEvent>,
    assets: HashMap<HandleId, Asset>,
}

impl AssetManager {
    pub fn new<P: AsRef<Path>>(path: P) -> Self {
        let (sender, receiver) = unbounded();
        Self {
            reader: get_default_reader(path),
            pending: Vec::default(),
            pending_fetched: Vec::default(),
            asset_id: 0,
            sender,
            receiver,
            assets: Default::default(),
        }
    }

    fn alloc_handle(&mut self) -> Handle {
        let id = self.asset_id;
        self.asset_id += 1;
        let drop_sender = self.sender.clone();
        Handle::new(id, drop_sender)
    }

    pub fn load<P: AsRef<Path>>(&mut self, asset_type: AssetType, path: P) -> Handle {
        let handle = self.alloc_handle();
        let task = PendingTask {
            handle: handle.clone(),
            path: path.as_ref().to_owned(),
            asset_type,
        };
        self.pending.push(task);
        if self.pending.len() > ASSET_COUNT_WARNING_THRESHOLD {
            log::warn!("Too many pending tasks");
        }
        handle
    }

    pub fn insert(&mut self, asset: Asset) -> Handle {
        let handle = self.alloc_handle();
        self.assets.insert(handle.id(), asset);
        if self.assets.len() > ASSET_COUNT_WARNING_THRESHOLD {
            log::warn!("Too many assets");
        }
        handle
    }

    pub fn load_texture<P: AsRef<Path>>(&mut self, path: P) -> Handle {
        self.load(AssetType::Texture, path)
    }

    pub fn load_font<P: AsRef<Path>>(&mut self, path: P) -> Handle {
        self.load(AssetType::Font, path)
    }

    pub fn load_bytes<P: AsRef<Path>>(&mut self, path: P) -> Handle {
        self.load(AssetType::Raw, path)
    }

    pub fn add_asset(&mut self, asset_type: AssetType, data: Vec<u8>) -> Result<Handle> {
        let handle = self.alloc_handle();
        self.register_asset(handle.clone(), asset_type, data)?;
        Ok(handle)
    }

    pub fn get_asset(&self, handle: &Handle) -> Option<&Asset> {
        self.assets.get(&handle.id())
    }

    pub fn remove_asset<P: AsRef<Path>>(&mut self, handle: &Handle) -> Option<Asset> {
        self.assets.remove(&handle.id())
    }

    pub fn has_pending_tasks(&self) -> bool {
        !self.pending.is_empty() || !self.pending_fetched.is_empty()
    }

    pub(crate) fn begin_fetch(&mut self) -> Result<Option<AssetFetchRequest>> {
        if self.pending.is_empty() {
            return Ok(None);
        }

        let tasks = mem::take(&mut self.pending)
            .into_iter()
            .map(|task| {
                let path = task
                    .path
                    .to_str()
                    .ok_or_else(|| anyhow!("Asset path is not valid UTF-8: {:?}", task.path))?
                    .to_owned();
                Ok(AssetFetchTask {
                    handle: task.handle,
                    asset_type: task.asset_type,
                    path,
                })
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(Some(AssetFetchRequest {
            reader: Arc::clone(&self.reader),
            tasks,
        }))
    }

    pub(crate) fn finish_fetch(
        &mut self,
        outputs: Vec<AssetFetchOutput>,
    ) -> Result<Vec<FetchedTask>> {
        for output in outputs {
            self.register_asset(output.handle, output.asset_type, output.bytes)?;
        }
        let mut tasks = mem::take(&mut self.pending_fetched);
        tasks.extend(self.collect_drop_tasks());
        Ok(tasks)
    }

    pub(crate) fn take_ready_tasks(&mut self) -> Vec<FetchedTask> {
        mem::take(&mut self.pending_fetched)
    }

    pub(crate) fn collect_drop_tasks(&mut self) -> Vec<FetchedTask> {
        let mut tasks = Vec::new();
        while let Ok(event) = self.receiver.try_recv() {
            let Some(asset) = self.assets.remove(&event.0) else {
                continue;
            };
            match asset.asset_type {
                AssetType::Raw => {}
                AssetType::Texture => {
                    tasks.push(FetchedTask::RemoveTexture { handle: event.0 });
                }
                AssetType::Font => {
                    tasks.push(FetchedTask::RemoveFont { handle: event.0 });
                }
            }
        }
        tasks
    }

    pub(crate) fn register_asset(
        &mut self,
        handle: Handle,
        asset_type: AssetType,
        data: Vec<u8>,
    ) -> Result<()> {
        match asset_type {
            AssetType::Raw => {
                self.assets.insert(
                    handle.id(),
                    Asset {
                        bytes: Some(data),
                        asset_type: AssetType::Raw,
                    },
                );
            }
            AssetType::Texture => {
                let im = image::load_from_memory(&data)?;
                let fetched_task = FetchedTask::CreateTexture {
                    handle: handle.clone(),
                    data: im,
                };
                self.pending_fetched.push(fetched_task);
                self.assets.insert(
                    handle.id(),
                    Asset {
                        bytes: None,
                        asset_type: AssetType::Texture,
                    },
                );
            }
            AssetType::Font => {
                let font = Font::from_bytes(data).ok_or(anyhow!("Failed to load font"))?;
                let fetched_task = FetchedTask::CreateFont {
                    handle: handle.clone(),
                    font,
                };
                self.pending_fetched.push(fetched_task);
                self.assets.insert(
                    handle.id(),
                    Asset {
                        bytes: None,
                        asset_type: AssetType::Font,
                    },
                );
            }
        }
        Ok(())
    }
}