bevy_modloader 0.1.2

use std::{mem::take, task::Poll};

use bevy_app::App;
use bevy_ecs::{
    event::Event,
    system::{Commands, Resource},
    world::{CommandQueue, World},
};
use bevy_tasks::{ComputeTaskPool, futures_lite::future::poll_fn};
use bevy_utils::HashMap;

use crate::mod_object::{GlobalModConfig, ModDependency, ModInfo, ModState};

use super::mod_object::Mod;

#[derive(Debug)]
enum ModLoaderStage {
    Parallel(Vec<Mod>),
    Single(Mod),
}

#[derive(Resource, Debug, Default)]
pub struct ModLoader {
    pub(crate) mods: Vec<Mod>,
    #[allow(private_interfaces)]
    pub(crate) mod_stages: Vec<ModLoaderStage>,
    pub(crate) global_state: ModState,
    pub(crate) config: GlobalModConfig,
}

/// Emitted when a mod is being unloaded while the program is running.
#[derive(Event, Clone, Debug)]
pub struct ModUnloading {
    pub mod_id: String,
    pub mod_state: ModState,
}

/// Emitted when a mod is being loaded while the program is running.
#[derive(Event, Clone, Debug)]
pub struct ModLoading {
    pub mod_id: String,
    pub mod_state: ModState,
}

impl ModLoader {
    /// Checks if the mod loader is ready to be ran.
    pub fn ready(&self) -> bool {
        self.mod_stages.len() > 0 || self.mod_stages.len() == 0 && self.mods.len() == 0
    }
    /// Adds a mod while the mod loader is still not ready.
    /// Returns an error if the mod loader is already loaded.
    pub fn add_mod(&mut self, m: Mod) -> Result<(), String> {
        if self.ready() {
            Err("Cannot add mod to already packed loader, you might want to unpack it.".to_string())
        } else {
            self.mods.push(m);
            Ok(())
        }
    }
    fn unpack_loader(&mut self) {
        if !self.ready() {
            return;
        }
        let mut unpacked = Vec::new();
        for stage in take(&mut self.mod_stages) {
            match stage {
                ModLoaderStage::Parallel(mods) => {
                    unpacked.extend(mods);
                }
                ModLoaderStage::Single(m) => {
                    unpacked.push(m);
                }
            }
        }
        self.mods = unpacked;
    }
    fn pack_loader(&mut self) -> Result<(), String> {
        if self.ready() {
            return Ok(());
        }
        let mut packed = Vec::new();
        let dependency_sorted;
        let mut mods_by_name = HashMap::new();
        take(&mut self.mods).into_iter().for_each(|m| {
            mods_by_name.insert(m.info.name.clone(), m);
        });

        // Create the dependency sort, check for circular dependencies and conflicts
        {
            let mut already_added = Vec::new();
            let mut adding = Vec::new();
            fn add(
                m: &ModInfo,
                adding: &mut Vec<String>,
                already_added: &mut Vec<String>,
                mods_by_name: &HashMap<String, Mod>,
            ) -> Result<(), String> {
                if already_added.contains(&m.name) {
                    return Ok(());
                }
                if adding.contains(&m.name) {
                    return Err(format!("Circular dependency detected for mod: {}", m.name));
                }
                adding.push(m.name.clone());

                for dep in m.dependencies.iter() {
                    match dep {
                        ModDependency::Conflict(name) => {
                            if let Some(mod_info) = mods_by_name.get(name) {
                                if already_added.contains(&mod_info.info.name) {
                                    return Err(format!(
                                        "Conflicting mod: {} with {}",
                                        m.name, mod_info.info.name
                                    ));
                                }
                            }
                        }
                        ModDependency::Required { name, version }
                        | ModDependency::Optional { name, version } => {
                            if let Some(mod_info) = mods_by_name.get(name) {
                                if !version.matches(&mod_info.info.version) {
                                    return Err(format!(
                                        "Mod {} requires version {} of {}, but found {}",
                                        m.name,
                                        version.to_string(),
                                        name,
                                        mod_info.info.version.to_string()
                                    ));
                                }
                                add(&mod_info.info, adding, already_added, mods_by_name)?;
                            } else {
                                return Err(format!("Missing required mod: {}", name));
                            }
                        }
                    }
                }
                already_added.push(adding.pop().unwrap());
                Ok(())
            }

            let mut collected_names = mods_by_name.keys().cloned().collect::<Vec<_>>();
            collected_names.sort_unstable();

            for info in collected_names
                .into_iter()
                .map(|name| &mods_by_name.get(&name).unwrap().info)
            {
                add(info, &mut adding, &mut already_added, &mods_by_name)?;
            }
            dependency_sorted = already_added;
        }
        // Create the packed stages
        dependency_sorted
            .into_iter()
            .map(|s| mods_by_name.remove(&s).unwrap())
            .for_each(|m| {
                if m.mod_loader.exclusive() {
                    packed.push(ModLoaderStage::Single(m));
                } else {
                    if let Some(last) = packed.last_mut() {
                        match last {
                            ModLoaderStage::Parallel(mods) => mods.push(m),
                            ModLoaderStage::Single(_) => {
                                packed.push(ModLoaderStage::Parallel(vec![m]));
                            }
                        }
                    } else {
                        packed.push(ModLoaderStage::Parallel(vec![m]));
                    }
                }
            });
        self.mod_stages = packed;
        Ok(())
    }

    /// Calls build on all mods that require building.
    /// Panics if the mod loader already has been built.
    pub fn build(&mut self, app: &mut App) {
        if self.ready() {
            return;
        }
        self.pack_loader().unwrap();
        for stage in self.mod_stages.iter_mut() {
            match stage {
                ModLoaderStage::Parallel(mods) => {
                    for m in mods.iter_mut() {
                        m.mod_loader.build(app);
                    }
                }
                ModLoaderStage::Single(m) => {
                    m.mod_loader.build(app);
                }
            }
        }
    }
    pub fn build_cleanup(&mut self, app: &mut App) {
        if !self.ready() {
            panic!("build_cleanup called on mod loader before build");
        }
        for stage in self.mod_stages.iter_mut() {
            match stage {
                ModLoaderStage::Parallel(mods) => {
                    for m in mods.iter_mut() {
                        m.mod_loader.build_cleanup(app);
                    }
                }
                ModLoaderStage::Single(m) => {
                    m.mod_loader.build_cleanup(app);
                }
            }
        }
    }

    /// Executes one mod stage.
    pub fn execute_stage(&mut self, state: ModState, world: &mut World) -> Result<(), String> {
        if !self.ready() {
            panic!("Mod loader not ready");
        }
        if state != self.global_state.consecutive() {
            panic!(
                "Mod loader expected state {:?} but got {:?}",
                self.global_state.consecutive(),
                state
            );
        }
        
        let task_pool = ComputeTaskPool::get();
        for stage in self.mod_stages.iter_mut() {
            match stage {
                ModLoaderStage::Parallel(mods) => {
                    let iter = mods
                        .iter_mut()
                        .map(|m| {
                            let c = &raw mut *self.config.entry_ref(&m.info.name).or_default();
                            (m, c)
                        })
                        .map(|(m, config)| (m, unsafe { &mut *config }));
                    let world_ref = &*world;
                    let results = task_pool.scope(move |s| {
                        for (m, config) in iter {
                            s.spawn(poll_fn(move |_| {
                                let mut command_buffer = CommandQueue::default();
                                let commands = Commands::new(&mut command_buffer, world_ref);
                                Poll::Ready(
                                    m.mod_loader
                                        .load(state, config, commands, world_ref)
                                        .map(move |_| command_buffer)
                                        .map_err(|e| {
                                            format!(
                                                "Mod {} {} failed: {}",
                                                m.info.name,
                                                m.info.version.to_string(),
                                                e
                                            )
                                        }),
                                )
                            }));
                        }
                    });
                    for i in results {
                        let mut command_buffer = i?;
                        command_buffer.apply(world);
                    }
                }
                ModLoaderStage::Single(m) => {
                    m.mod_loader.load_exclusive(
                        state,
                        self.config.entry_ref(&m.info.name).or_default(),
                        world,
                    )?;
                }
            }
        }
        self.global_state = state;
        Ok(())
    }
}