engine/rpg_components/src/items/item_system.rs

use anyhow::Result;
use assetpath::AssetPath;
use engine::prelude::*;

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use image::{imageops::FilterType, DynamicImage, ImageBuffer, Pixel, Rgba, RgbaImage};

use crate::components::attributes::{Attribute, Attributes};
use crate::components::inventory::Storable;
use crate::components::statistic_types::{
    AgilityStatisticTypes, IntelligenceStatisticTypes, StatisticType, StrengthStatisticTypes,
};
use crate::config::abilities::AbilitySettings;
use crate::config::attributes::AttributeSettings;
use crate::config::items::ItemSettings;

use super::ability_addon::{AbilityAddon, AbilityAddonTypes};
use super::ability_book::{Ability, AbilityBook};
use super::{Item, ItemAffix, ItemSlots, Jewel, Rarities};

#[derive(Debug, Clone, PartialEq)]
pub enum Loot<A: Ability> {
    Item(Item),
    AbilityBook(AbilityBook<A>),
    AbilityAddOn(AbilityAddon),
    Jewel(Jewel),
}

impl<A: Ability> Loot<A> {
    pub fn storable(&self) -> &dyn Storable {
        match self {
            Loot::Item(item) => item,
            Loot::AbilityBook(book) => book,
            Loot::AbilityAddOn(addon) => addon,
            Loot::Jewel(jewel) => jewel,
        }
    }
}

pub struct ItemSystem<A: Ability> {
    pub item_settings: ItemSettings,
    ability_settings: AbilitySettings,

    item_icon_combinations: HashMap<(Rarities, ItemSlots), Arc<Image>>,
    ability_icon_combinations: HashMap<(Rarities, String), Arc<Image>>,
    addon_icon_combinations: HashMap<(Rarities, AbilityAddonTypes), Arc<Image>>,
    jewel_icon_combinations: HashMap<(Rarities, u32, Attribute), Arc<Image>>,

    abilities: Vec<A>,
}

impl<A: Ability> ItemSystem<A> {
    pub fn new(
        engine: &Arc<Engine>,
        item_settings: &ItemSettings,
        ability_settings: &AbilitySettings,
        attribute_settings: &AttributeSettings,
        ability_directory: &AssetPath,
    ) -> Result<Self> {
        // verify that drop chances sum up to 1.0
        item_settings.rarity_drop_rates.verify();

        let abilities = search_dir_recursively(&ability_directory.full_path(), ".abil")?
            .into_iter()
            .map(|path| A::create(path))
            .collect::<Result<Vec<A>>>()?;

        let (
            item_icon_combinations,
            ability_icon_combinations,
            addon_icon_combinations,
            jewel_icon_combinations,
        ) = {
            // images for item slots
            let mut slot_images = Vec::new();

            for slot in ItemSlots::iter() {
                let path = slot.get_path(item_settings);

                if !path.is_empty() {
                    slot_images.push((*slot, Self::dyn_image(path)?));
                }
            }

            // image for ability book
            let mut ability_images = Vec::new();

            for loader in abilities.iter() {
                let path = loader.icon_path();

                if !path.is_empty() {
                    ability_images.push((loader.name().to_string(), Self::dyn_image(path)?));
                }
            }

            // images for ability addons
            let mut ability_addon_images = Vec::new();

            for addon in AbilityAddonTypes::iter() {
                let path = addon.get_path(ability_settings);

                if !path.is_empty() {
                    ability_addon_images.push((*addon, Self::dyn_image(path)?));
                }
            }

            // images for jewels
            let mut jewel_images = Vec::new();
            jewel_images.push((0, Self::dyn_image(&item_settings.jewel_paths.first)?));
            jewel_images.push((1, Self::dyn_image(&item_settings.jewel_paths.second)?));
            jewel_images.push((2, Self::dyn_image(&item_settings.jewel_paths.third)?));
            jewel_images.push((3, Self::dyn_image(&item_settings.jewel_paths.fourth)?));

            // resize all icons to have the same size
            let (
                item_icons,
                base_background,
                ability_icons,
                ability_addon_icons,
                addon_background,
                jewel_icons,
            ) = Self::equalize_images(
                slot_images,
                Self::dyn_image(&item_settings.icon_paths.background)?,
                ability_images,
                ability_addon_images,
                Self::dyn_image(&ability_settings.icons.addon_background)?,
                jewel_images,
                item_settings,
            );

            // create item icon combination for every slot and rarity
            let mut item_icon_combinations = HashMap::new();

            for (slot, icon) in item_icons.iter() {
                for rarity in Rarities::iter() {
                    let final_image = Self::blend_background(
                        icon,
                        &rarity.apply_color::<A>(
                            &base_background,
                            &item_settings.rarity_color_settings,
                        ),
                    );

                    let (width, height) = final_image.dimensions();

                    item_icon_combinations.insert(
                        (*rarity, *slot),
                        Self::create_icon(
                            engine.device(),
                            engine.queue(),
                            width,
                            height,
                            final_image.into_raw(),
                        )?,
                    );
                }
            }

            // create book icon combination for every rarity
            let mut ability_icon_combinations = HashMap::new();

            for (name, image) in ability_icons.iter() {
                for rarity in Rarities::iter() {
                    let final_image = Self::blend_background(
                        &image,
                        &rarity.apply_color::<A>(
                            &base_background,
                            &item_settings.rarity_color_settings,
                        ),
                    );

                    let (width, height) = final_image.dimensions();

                    ability_icon_combinations.insert(
                        (*rarity, name.clone()),
                        Self::create_icon(
                            engine.device(),
                            engine.queue(),
                            width,
                            height,
                            final_image.into_raw(),
                        )?,
                    );
                }
            }

            // create addon icon combinatiosn for every rarity and addon type
            let mut addon_icon_combinations = HashMap::new();

            for (addon_type, icon) in ability_addon_icons.iter() {
                for rarity in Rarities::iter() {
                    let final_image = Self::blend_apply(
                        &rarity.apply_color::<A>(
                            &addon_background,
                            &item_settings.rarity_color_settings,
                        ),
                        icon,
                        |rarity_color, icon_color| {
                            // let (r1, g1, b1, a1) = rarity_color.channels4();
                            let c2 = icon_color.channels();

                            if c2[3] == 0 {
                                rarity_color
                            } else {
                                Rgba([c2[0], c2[1], c2[2], 255])
                            }
                        },
                    );

                    let (width, height) = final_image.dimensions();

                    addon_icon_combinations.insert(
                        (*rarity, *addon_type),
                        Self::create_icon(
                            engine.device(),
                            engine.queue(),
                            width,
                            height,
                            final_image.into_raw(),
                        )?,
                    );
                }
            }

            // create jewel icons for every rarity
            let mut jewel_icon_combinations = HashMap::new();

            for (level, icon) in jewel_icons.iter() {
                for attribute in Attribute::iter() {
                    let attribute_icon = attribute
                        .apply_color::<A>(&icon, &attribute_settings.attribute_color_settings);

                    for rarity in Rarities::iter() {
                        let final_image = Self::blend_background(
                            &attribute_icon,
                            &rarity.apply_color::<A>(
                                &base_background,
                                &item_settings.rarity_color_settings,
                            ),
                        );

                        let (width, height) = final_image.dimensions();

                        jewel_icon_combinations.insert(
                            (*rarity, *level + 1, *attribute),
                            Self::create_icon(
                                engine.device(),
                                engine.queue(),
                                width,
                                height,
                                final_image.into_raw(),
                            )?,
                        );
                    }
                }
            }

            (
                item_icon_combinations,
                ability_icon_combinations,
                addon_icon_combinations,
                jewel_icon_combinations,
            )
        };

        Ok(ItemSystem {
            item_settings: item_settings.clone(),
            ability_settings: ability_settings.clone(),

            item_icon_combinations,
            ability_icon_combinations,
            addon_icon_combinations,
            jewel_icon_combinations,

            abilities,
        })
    }

    #[inline]
    fn create_icon(
        device: &Arc<Device>,
        queue: &Arc<Mutex<Queue>>,
        width: u32,
        height: u32,
        data: Vec<u8>,
    ) -> Result<Arc<Image>> {
        Image::from_raw(data, width, height)
            .format(VK_FORMAT_R8G8B8A8_UNORM)
            .attach_sampler(Sampler::nearest_sampler().build(device)?)
            .build(device, queue)
    }

    /// only here for debugging
    pub fn get_legendary_random_loot(&self, level: u32) -> Loot<A> {
        // decide which type of loot gets dropped
        let loot_type_p = Random::range_f32(0.0, 1.0);
        let rarity = Rarities::Legendary;

        // decide if an item gets dropped
        if loot_type_p >= 0.0 && loot_type_p <= self.item_settings.type_drop_rates.item {
            Loot::Item(self.random_item(level, rarity))
        }
        // decide if an ability addon gets dropped
        else if loot_type_p > self.item_settings.type_drop_rates.item
            && loot_type_p
                <= self.item_settings.type_drop_rates.item
                    + self.item_settings.type_drop_rates.ability_addon
        {
            Loot::AbilityAddOn(self.random_ability_addon(rarity))
        }
        // drop an ability book
        else {
            if Coin::flip(0.5) {
                Loot::AbilityBook(self.random_ability_book(rarity))
            } else {
                Loot::Jewel(self.random_jewel(level, rarity))
            }
        }
    }

    pub fn get_random_loot(
        &self,
        level: u32,
        drop_chance_multiplier: Option<f32>,
    ) -> Option<Loot<A>> {
        let drop_chance = match drop_chance_multiplier {
            Some(multiplier) => multiplier * self.item_settings.general.drop_chance,
            None => self.item_settings.general.drop_chance,
        };

        // decide if loot gets dropped at all
        if !Coin::flip(drop_chance) {
            return None;
        }

        // decide which type of loot gets dropped
        let loot_type_p = Random::range_f32(0.0, 1.0);

        let rarity = Rarities::random(&self.item_settings, level);

        // decide if an item gets dropped
        if loot_type_p >= 0.0 && loot_type_p <= self.item_settings.type_drop_rates.item {
            Some(Loot::Item(self.random_item(level, rarity)))
        }
        // decide if an ability addon gets dropped
        else if loot_type_p > self.item_settings.type_drop_rates.item
            && loot_type_p
                <= self.item_settings.type_drop_rates.item
                    + self.item_settings.type_drop_rates.ability_addon
        {
            Some(Loot::AbilityAddOn(self.random_ability_addon(rarity)))
        }
        // drop an ability book
        else {
            if Coin::flip(0.5) {
                Some(Loot::AbilityBook(self.random_ability_book(rarity)))
            } else {
                Some(Loot::Jewel(self.random_jewel(level, rarity)))
            }
        }
    }

    pub fn random_ability_book(&self, rarity: Rarities) -> AbilityBook<A> {
        let ability = self.random_ability();
        let ability_icon = self.ability_icon(ability.name(), rarity);

        AbilityBook::new(ability, ability_icon, rarity, &self.ability_settings)
    }

    pub fn ability_addon(
        &self,
        rarity: Rarities,
        mut addon_type: AbilityAddonTypes,
    ) -> AbilityAddon {
        let icon = self.addon_icon(rarity, addon_type);

        addon_type.apply_rarity(rarity, &self.ability_settings);

        AbilityAddon::new(rarity, addon_type, icon)
    }

    pub fn random_ability_addon(&self, rarity: Rarities) -> AbilityAddon {
        self.ability_addon(rarity, AbilityAddonTypes::random())
    }

    pub fn random_item(&self, level: u32, rarity: Rarities) -> Item {
        let slot = ItemSlots::random();

        let mut item = Item::empty(rarity, slot, level, self.item_icon(rarity, slot));

        item.randomize_attributes(&self.item_settings);
        item.randomize_stats(&self.item_settings);

        item
    }

    pub fn random_jewel(&self, level: u32, rarity: Rarities) -> Jewel {
        let attribute = Attribute::random();

        let stat = match attribute {
            Attribute::Agility => AgilityStatisticTypes::random()
                .apply_for_jewel(level, rarity, &self.item_settings)
                .into(),
            Attribute::Intelligence => IntelligenceStatisticTypes::random()
                .apply_for_jewel(level, rarity, &self.item_settings)
                .into(),
            Attribute::Strength => StrengthStatisticTypes::random()
                .apply_for_jewel(level, rarity, &self.item_settings)
                .into(),
        };

        self.jewel(rarity, level, attribute, stat)
    }

    pub fn item(
        &self,
        rarity: Rarities,
        slot: ItemSlots,
        level: u32,
        attributes: Attributes,
        mut affixes: Vec<ItemAffix>,
    ) -> Item {
        for affix in affixes.iter_mut() {
            match affix {
                ItemAffix::Socket(_) => (),
                ItemAffix::Stat(stat) => stat.apply_for_item(&attributes, &self.item_settings),
            }
        }

        Item {
            rarity,
            slot,
            level,
            attributes,
            affixes,

            icon: self.item_icon(rarity, slot),
        }
    }

    pub fn addon(&self, rarity: Rarities, addon_type: AbilityAddonTypes) -> AbilityAddon {
        AbilityAddon::load(
            addon_type,
            rarity,
            self.addon_icon(rarity, addon_type),
            &self.ability_settings,
        )
    }

    pub fn ability_book(
        &self,
        ability_name: &str,
        rarity: Rarities,
        addons: Vec<Option<AbilityAddon>>,
        level: u32,
    ) -> AbilityBook<A> {
        AbilityBook::load(
            self.find_ability(ability_name),
            self.ability_icon(ability_name, rarity),
            rarity,
            addons,
            &self.ability_settings,
            level,
        )
    }

    pub fn jewel(
        &self,
        rarity: Rarities,
        level: u32,
        attribute: Attribute,
        stat: StatisticType,
    ) -> Jewel {
        Jewel {
            rarity,
            level,
            attribute,

            stat,

            icon: Some(self.jewel_icon(rarity, level, attribute)),
        }
    }

    fn dyn_image(path: &AssetPath) -> Result<DynamicImage> {
        Ok(image::open(&path.full_path())?)
    }

    #[inline]
    pub fn item_icon(&self, rarity: Rarities, slot: ItemSlots) -> Arc<Image> {
        self.item_icon_combinations
            .get(&(rarity, slot))
            .unwrap_or_else(|| {
                panic!(
                    "no icon for combination ({:?} {:?}) present\navailable: {:#?}",
                    rarity,
                    slot,
                    self.item_icon_combinations.keys()
                )
            })
            .clone()
    }

    #[inline]
    pub fn addon_icon(&self, rarity: Rarities, addon_type: AbilityAddonTypes) -> Arc<Image> {
        let addon_type = addon_type.into_zero();

        self.addon_icon_combinations
            .get(&(rarity, addon_type))
            .unwrap_or_else(|| {
                panic!(
                    "no addon icon present for rarity ({}) and addon type ({:?})\navailable: {:#?}",
                    rarity,
                    addon_type,
                    self.addon_icon_combinations.keys()
                )
            })
            .clone()
    }

    #[inline]
    pub fn ability_icon(&self, name: impl ToString, rarity: Rarities) -> Arc<Image> {
        self.ability_icon_combinations
            .get(&(rarity, name.to_string()))
            .unwrap_or_else(|| {
                panic!(
                    "no book icon present for rarity ({})\navailable: {:#?}",
                    rarity,
                    self.ability_icon_combinations.keys()
                )
            })
            .clone()
    }

    #[inline]
    pub fn jewel_icon(&self, rarity: Rarities, level: u32, attribute: Attribute) -> Arc<Image> {
        self.jewel_icon_combinations
            .get(&(rarity, level, attribute))
            .unwrap_or_else(|| {
                panic!(
                    "no jewel icon present for rarity ({})\navailable: {:#?}",
                    rarity,
                    self.jewel_icon_combinations.keys()
                )
            })
            .clone()
    }

    #[inline]
    pub fn find_ability(&self, name: &str) -> A {
        self.abilities
            .iter()
            .find(|a| a.name() == name)
            .unwrap_or_else(|| panic!("no ability with name ({}) found", name))
            .clone()
    }

    #[inline]
    pub fn random_ability(&self) -> A {
        let n = Random::range(0, self.abilities.len() as u32);
        self.abilities[n as usize].clone()
    }

    fn equalize_images(
        item_icons: Vec<(ItemSlots, DynamicImage)>,
        background: DynamicImage,
        ability_images: Vec<(String, DynamicImage)>,
        addon_icons: Vec<(AbilityAddonTypes, DynamicImage)>,
        addon_background: DynamicImage,
        jewel_images: Vec<(u32, DynamicImage)>,
        item_settings: &ItemSettings,
    ) -> (
        Vec<(ItemSlots, RgbaImage)>,
        RgbaImage,
        Vec<(String, RgbaImage)>,
        Vec<(AbilityAddonTypes, RgbaImage)>,
        RgbaImage,
        Vec<(u32, RgbaImage)>,
    ) {
        let rgba_background = background
            .resize(
                item_settings.general.icon_target_width,
                item_settings.general.icon_target_height,
                FilterType::Gaussian,
            )
            .to_rgba8();

        let add_on_background = addon_background
            .resize(
                item_settings.general.icon_target_width,
                item_settings.general.icon_target_height,
                FilterType::Gaussian,
            )
            .to_rgba8();

        let icon_results = item_icons
            .into_iter()
            .map(|(slot, image)| {
                (
                    slot,
                    image
                        .resize(
                            item_settings.general.icon_target_width,
                            item_settings.general.icon_target_height,
                            FilterType::Gaussian,
                        )
                        .to_rgba8(),
                )
            })
            .collect();

        let ability_icon = ability_images
            .into_iter()
            .map(|(name, image)| {
                (
                    name,
                    image
                        .resize(
                            item_settings.general.icon_target_width,
                            item_settings.general.icon_target_height,
                            FilterType::Gaussian,
                        )
                        .to_rgba8(),
                )
            })
            .collect();

        let ability_addons = addon_icons
            .into_iter()
            .map(|(addon, image)| {
                (
                    addon,
                    image
                        .resize(
                            item_settings.general.icon_target_width,
                            item_settings.general.icon_target_height,
                            FilterType::Gaussian,
                        )
                        .to_rgba8(),
                )
            })
            .collect();

        let jewel_results = jewel_images
            .into_iter()
            .map(|(n, image)| {
                (
                    n,
                    image
                        .resize(
                            item_settings.general.icon_target_width,
                            item_settings.general.icon_target_height,
                            FilterType::Gaussian,
                        )
                        .to_rgba8(),
                )
            })
            .collect();

        (
            icon_results,
            rgba_background,
            ability_icon,
            ability_addons,
            add_on_background,
            jewel_results,
        )
    }

    pub fn blend_background(i1: &RgbaImage, i2: &RgbaImage) -> RgbaImage {
        let (width, height) = i1.dimensions();
        let mut target = ImageBuffer::new(width, height);

        for ((s1, s2), t) in i1.pixels().zip(i2.pixels()).zip(target.pixels_mut()) {
            let mut source_clone = *s1;
            source_clone.blend(s2);

            *t = source_clone;
        }

        target
    }

    pub fn blend_apply<F>(i1: &RgbaImage, i2: &RgbaImage, mut apply: F) -> RgbaImage
    where
        F: FnMut(Rgba<u8>, &Rgba<u8>) -> Rgba<u8> + Copy,
    {
        let (width, height) = i1.dimensions();
        let mut target = ImageBuffer::new(width, height);

        for ((s1, s2), t) in i1.pixels().zip(i2.pixels()).zip(target.pixels_mut()) {
            *t = apply(*s1, s2);
        }

        target
    }

    #[inline]

    pub fn apply_color(base_image: &RgbaImage, result_color: Color) -> RgbaImage {
        let mut result_image = base_image.clone();

        for color in result_image.pixels_mut() {
            let c = color.channels();
            *color = Self::blend_mul((c[0], c[1], c[2], c[3]), result_color.into());
        }

        result_image
    }

    #[inline]

    fn blend_mul((r, g, b, a): (u8, u8, u8, u8), c: [u8; 3]) -> Rgba<u8> {
        Rgba([
            Self::mul(r, c[0]),
            Self::mul(g, c[1]),
            Self::mul(b, c[2]),
            a,
        ])
    }

    #[inline]

    fn mul(l: u8, r: u8) -> u8 {
        let lhs = l as f32 / 255.0;
        let rhs = r as f32 / 255.0;

        (lhs * rhs * 255.0) as u8
    }
}

#[test]
fn validate_item_icons() {
    let item_settings = ItemSettings::default();

    let mut slots = Vec::new();

    for slot in ItemSlots::iter() {
        let path = slot.get_path(&item_settings);

        slots.push((*slot, path));
    }

    let mut item_icon_combinations = std::collections::HashMap::new();

    for (slot, path) in slots.iter() {
        for rarity in Rarities::iter() {
            item_icon_combinations.insert((*slot, *rarity), (*path).clone());
        }
    }

    for slot in ItemSlots::iter() {
        for rarity in Rarities::iter() {
            item_icon_combinations
                .get(&(*slot, *rarity))
                .unwrap_or_else(|| panic!("item not present for {:?} + {}", slot, rarity));
        }
    }
}

#[test]
fn validate_addon_icons() {
    let ability_settings = AbilitySettings::default();

    let mut addon_types = Vec::new();

    for addon_type in AbilityAddonTypes::iter() {
        let path = addon_type.get_path(&ability_settings);

        addon_types.push((*addon_type, path));
    }

    let mut addon_icon_combinations = std::collections::HashMap::new();

    for (addon_type, path) in addon_types.iter() {
        for rarity in Rarities::iter() {
            addon_icon_combinations.insert((*rarity, *addon_type), (*path).clone());
        }
    }

    for addon_type in AbilityAddonTypes::iter() {
        for rarity in Rarities::iter() {
            addon_icon_combinations
                .get(&(*rarity, *addon_type))
                .unwrap_or_else(|| panic!("item not present for {:?} + {}", addon_type, rarity));
        }
    }
}