/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! # Tiny ECS
//!
//! The intention of this crate is that a basic ECS is provided, where
//! you will be required to exercise a little additional control. This is
//! somewhat due to some limitations, and also due to trying to maintain
//! as little overhead as possible - this means no unneccesary copies/clones,
//! and in some use cases you will need to use `downcast_ref` or `downcast_mut`.
//!
//! The basis of this ECS is the use of `bitmasks`. Each entity ID is in
//! practice an internal index number in to an array which contains bitmasks.
//! The bitmasks themselves keep track of what components the entity has.
//! For the most part, bitmasks are handled for you, and some helper methods
//! are available to hide their use, but there are also methods to get the
//! bitmask for any ID if you are inclined to do some manual management.
//!
//! # Examples
//!
//! **Init with a capacity.**
//!
//! This is good to do if you know the size required
//! as it will prevent many reallocs/moves as data is added. This affects
//! both the entity and partmaps allocs (they will be equal in size).
//!
//! ```
//! use tiny_ecs::Entities;
//!
//! let mut entities = Entities::new(Some(1000));
//! ```
//!
//! **Demonstrating use**
//!
//! ```
//! use tiny_ecs::{Entities, PartMap};
//! use std::any::Any;
//! use std::collections::HashMap;
//! use std::cell::{Ref, RefMut};
//!
//! struct Vector1 { x: i32 }
//!
//! struct Vector2 { x: i32, y: i32 }
//!
//! struct Vector3 { x: i32, y: i32, z: i32 }
//!
//! let mut entities = Entities::new(Some(10));
//! let entity_1 = entities.get_free_slot().unwrap();
//!
//! assert!(entities
//!         .add_part(entity_1,
//!                   Vector1 { x: 42 }).is_ok());
//! assert!(entities
//!         .add_part(entity_1,
//!                   Vector3 { x: 3, y: 10, z: -12 }).is_ok());
//!
//! let entity_2 = entities.get_free_slot().unwrap();
//! assert!(entities
//!         .add_part(entity_2,
//!                   Vector2 { x: 66, y: 6 }).is_ok());
//! assert!(entities
//!         .add_part(entity_2,
//!                   Vector1 { x: 6 }).is_ok());
//!
//! if entities.entity_contains::<Vector3>(entity_1) {
//!     let mut partmap = entities.get_pmap_mut::<Vector3>().unwrap().unwrap();
//!     let mut part: &Vector3 = partmap.get_part_mut(entity_1).unwrap();
//!     assert_eq!(part.z, -12);
//! }
//!
//! if entities.entity_contains::<Vector1>(entity_1) {
//!     assert!(entities.rm_part::<Vector1>(entity_1));
//! }
//! assert_eq!(entities.entity_contains::<Vector1>(entity_1), false);
//!
//! fn some_system(mut partmap: RefMut<PartMap>) {
//!     for (k, v) in partmap.iter_mut() {
//!         let part = v.downcast_mut::<Vector1>().unwrap();
//!         part.x += 1;
//!         assert!(part.x > *k as i32);
//!     }
//! }
//! some_system(entities.get_pmap_mut::<Vector1>().unwrap().unwrap());
//!
//! fn second_system(active: &[u32], mut v1_map: RefMut<PartMap>) {
//!     for id in active {
//!         if let Ok(part) = v1_map.get_part_mut::<Vector1>(*id) {
//!              part.x = 42;
//!         }
//!     }
//! }
//! second_system(&[0, 1, 2], entities.get_pmap_mut::<Vector1>().unwrap().unwrap());
//!
//! // Or a system can use a list of entities you keep track of
//! fn other_system(active_ents: &[u32], entities: &mut Entities) {
//!     let mut partmap = entities.get_pmap_mut::<Vector1>()
//!                                 .expect("No part found for entity")
//!                                 .expect("Mutable borrow failed");
//!     for id in active_ents {
//!         if entities.entity_contains::<Vector1>(*id) {
//!             let part: &mut Vector1 = partmap.get_part_mut(*id).unwrap();
//!             part.x = 42;
//!             assert_ne!(part.x, 43);
//!             assert_eq!(part.x, 42);
//!         }
//!     }
//! }
//! other_system(&[0, 1, 2], &mut entities);
//! ```

use std::any::{Any, TypeId};
use std::cell::{BorrowError, BorrowMutError, Ref, RefCell, RefMut};
use std::collections::HashMap;
use std::collections::hash_map::{Iter, IterMut};

/// Used to fill the initial mask list, and replace deleted entities
pub const EMPTY: u32 = 0;

/// `PartMap` is a container type for the parts used by entities
pub struct PartMap {
    map: HashMap<u32, Box<Any>>,
}

impl PartMap {
    fn new(init_size: Option<usize>) -> PartMap {
        PartMap {
            map: HashMap::with_capacity(init_size.unwrap_or(0)),
        }
    }

    fn insert<T: 'static>(&mut self, id: u32, part: T) {
        self.map.insert(id, Box::new(part));
    }

    fn remove(&mut self, id: u32) {
        self.map.remove(&id);
    }

    pub fn get_part_ref<T: 'static>(&self, id: u32) -> Result<&T, &'static str> {
        if let Some(part) = self.map.get(&id) {
            let part = part.downcast_ref::<T>().unwrap();
            return Ok(part);
        }
        Err("could find requested entity part")
    }

    /// Getting a reference to an entity part requires the entity ID and
    /// type signature. The type signature is required so that the internal
    /// downcast can be done for you.
    ///
    /// # Example
    ///
    /// ```
    /// # use tiny_ecs::Entities;
    /// #[derive(Debug, PartialEq)]
    /// struct Test1 {}
    ///
    /// # let mut entities = Entities::new(Some(10));
    /// # let entity_1 = entities.get_free_slot().unwrap();
    /// # let _ = entities.add_part(entity_1, Test1 {}).unwrap();
    /// let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
    /// let part = partmap.get_part_ref::<Test1>(entity_1).unwrap();
    /// assert_eq!(part, &Test1 {});
    /// // or
    /// let part: &Test1 = partmap.get_part_ref(entity_1).unwrap();
    /// assert_eq!(part, &Test1 {});
    /// ```
    pub fn get_part_mut<T: 'static>(&mut self, id: u32) -> Result<&mut T, &'static str> {
        if let Some(part) = self.map.get_mut(&id) {
            let part = part.downcast_mut::<T>().unwrap();
            return Ok(part);
        }
        Err("could find requested entity part")
    }

    /// Returns an immutable iterator over the part map
    ///
    /// # Example
    ///
    /// ```
    /// # use tiny_ecs::Entities;
    /// struct Test1 { x: u32 }
    /// let mut entities = Entities::new(Some(10));
    ///
    /// let mut ids = Vec::new();
    /// let id = entities.get_free_slot().unwrap();
    /// assert!(entities.add_part(id, Test1 { x: id }).is_ok());
    /// ids.push(id);
    ///
    /// let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
    /// for (k, v) in partmap.iter_ref() {
    ///     // You must downcast the `Any` type for it to be usable
    ///     let part = v.downcast_ref::<Test1>().unwrap();
    ///     assert!(part.x == *k);
    /// }
    ///```
    pub fn iter_ref(&self) -> Iter<u32, Box<Any>> {
        return self.map.iter()
    }

    /// Returns a mutable iterator over the part map
    pub fn iter_mut(&mut self) -> IterMut<u32, Box<Any>> {
        return self.map.iter_mut()
    }
}

/// This is the root of the ECS implementation
pub struct Entities {
    entity_masks: Vec<u32>,
    parts: HashMap<u32, RefCell<PartMap>>,
    next_free_entity: u32,
    type_masks: HashMap<TypeId, u32>,
    next_type_bitshift: u32,
}

impl Default for Entities {
    /// Create a new `Entities` struct with no pre-allocated memory for maps
    fn default() -> Self {
        Entities::new(None)
    }
}

impl Entities {
    /// `init_size` is a required arg. The value given here is used to
    /// pre-allocate memory of `init_size` num of elements for both
    /// the entities themselves, and the maps containing their parts.
    pub fn new(init_size: Option<usize>) -> Entities {
        Entities {
            entity_masks: vec![EMPTY; init_size.unwrap_or(0)],
            parts: HashMap::with_capacity(init_size.unwrap_or(0)),
            next_free_entity: 0,
            type_masks: HashMap::new(),
            next_type_bitshift: 1,
        }
    }

    /// Find the first `EMPTY` ID number to use
    ///
    /// As entities are only created by actually inserting new parts in to
    /// the Entity structure, this should be called to find the first
    /// available entity slot. Slot states are determined by the mask
    /// it holds.
    pub fn get_free_slot(&mut self) -> Result<u32, &'static str> {
        if self.next_free_entity >= self.entity_masks.len() as u32 {
            self.entity_masks.push(EMPTY);
        }
        for index in 0..=self.next_free_entity {
            if self.entity_masks[index as usize] == EMPTY {
                if index >= self.next_free_entity {
                    self.next_free_entity += 1;
                }
                return Ok(index);
            }
        }
        Err("no free entity slots")
    }

    /// Returns the mask of the requested entity enabling you to manually
    /// check composition using bitmasks.
    pub fn get_entity_mask(&self, id: u32) -> u32 {
        self.entity_masks[id as usize]
    }

    /// Returns the mask associated with the requested type.
    pub fn get_type_mask<T: 'static>(&self) -> Result<u32, &'static str> {
        if let Some(mask) = self.type_masks.get(&TypeId::of::<T>()) {
            return Ok(*mask);
        }
        Err("requested non-existant mask for type")
    }

    pub fn entity_contains<T: 'static>(&self, id: u32) -> bool {
        if let Some(entity_mask) = self.entity_masks.get(id as usize) {
            if let Some(type_mask) = self.type_masks.get(&TypeId::of::<T>()) {
                if entity_mask & type_mask == *type_mask {
                    return true;
                }
            }
        }
        false
    }

    /// Adding a part requires a valid slot along with the initialised data
    /// to use with the entity. Effectively creates the entity if the slot is
    /// currently empty.
    ///
    /// A bitmask is created internally for each data type added (only one per type).
    ///
    /// # Example
    ///
    /// ```
    /// # use tiny_ecs::Entities;
    /// struct Test1 {}
    /// const TEST1: u32 = 1 << 2;
    ///
    /// # let mut entities = Entities::new(Some(10));
    /// let entity_1 = entities.get_free_slot().unwrap();
    /// assert!(entities.add_part(entity_1, Test1 {}).is_ok());
    /// ```
    pub fn add_part<P>(&mut self, id: u32, part: P) -> Result<(), &'static str>
    where
        P: Any,
    {
        let tid = TypeId::of::<P>();
        if !self.type_masks.contains_key(&tid) {
            let mask = self.next_type_bitshift << 1;
            self.type_masks.insert(tid, mask);
            self.next_type_bitshift += 1;
        }
        let type_mask = self.type_masks[&tid];
        // add HashMap for these parts if not exist
        self.parts
            .entry(type_mask)
            .or_insert(RefCell::new(PartMap::new(None)));

        if let Some(col) = self.parts.get_mut(&type_mask) {
            col.borrow_mut().insert(id, part);
            let old = self.entity_masks[id as usize];
            self.entity_masks[id as usize] = old | type_mask;
            return Ok(());
        }
        Err("could not add entity part")
    }

    /// Remove an entities part. If no parts are left after part removal then
    /// the entity is considered deleted and an `EMPTY` mask is inserted in
    /// its place.
    ///
    /// Removal requires the ID of the entity and the parts type signature.
    ///
    /// # Example
    ///
    /// ```
    /// # use tiny_ecs::{Entities, EMPTY};
    /// #[derive(Debug, PartialEq)]
    /// struct Test1 {}
    ///
    /// # let mut entities = Entities::new(Some(10));
    /// # let entity_1 = entities.get_free_slot().unwrap();
    /// # let _ = entities.add_part(entity_1, Test1 {}).unwrap();
    /// assert!(entities.rm_part::<Test1>(entity_1));
    /// assert_eq!(entities.get_entity_mask(entity_1), EMPTY);
    /// ```
    pub fn rm_part<T: 'static>(&mut self, id: u32) -> bool {
        let type_mask = self.type_masks[&TypeId::of::<T>()];
        if let Some(map) = self.parts.get_mut(&type_mask) {
            map.borrow_mut().remove(id);
            self.entity_masks[id as usize] ^= type_mask;
            return true;
        }
        false
    }

    /// Will remove all parts associated with an entity ID and replace
    /// with an `EMPTY` mask.
    pub fn rm_entity(&mut self, id: u32) {
        for (_, part_map) in &mut self.parts {
            part_map.borrow_mut().remove(id);
        }
        self.entity_masks[id as usize] = EMPTY;
    }

    /// Get a plain reference to the selected entity part map
    ///
    /// You may have multiple immutable references to the requested `PartMap`
    /// **type** but no mutable references if the same **typed** `PartMap`
    /// is currently referenced.
    ///
    /// - `Option` is whether or not there is a part of `<T>` for that entity.
    /// - `Result` covers if the map was able to be borrowed or not.
    /// - Borrowing is checked at runtime.
    ///
    /// # Example
    /// ```
    /// # use tiny_ecs::Entities;
    /// # #[derive(Debug, PartialEq)]
    /// struct Test1 { x: u32 }
    ///
    /// let mut entities = Entities::new(Some(10));
    /// // entity_1 is functionally the entity ID
    /// let entity_1 = entities.get_free_slot().unwrap();
    /// # assert!(entities.add_part(entity_1, Test1 { x: 666 }).is_ok());
    /// let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
    /// let part = partmap.get_part_ref::<Test1>(entity_1).unwrap();
    /// ```
    pub fn get_pmap_ref<T>(&self) -> Option<Result<Ref<PartMap>, BorrowError>>
    where
        T: 'static,
    {
        if let Some(type_mask) = self.type_masks.get(&TypeId::of::<T>()) {
            return Some(self.parts.get(&type_mask).unwrap().try_borrow());
        }
        None
    }

    /// Get a mutable reference to the selected entity part map
    ///
    /// You may have only one mutable reference to the requested `PartMap`
    /// **type** and no immutable references. You can however, have multiple
    /// mutable references to different **types** of `PartMap`
    ///
    /// - `Option` is whether or not there is a part of `<T>` for that entity.
    /// - `Result` covers if the map was able to be borrowed mutably or not.
    /// - Borrowing is checked at runtime.
    ///
    /// # Example
    /// ```
    /// # use tiny_ecs::Entities;
    /// # #[derive(Debug, PartialEq)]
    /// struct Test1 { x: u32 }
    ///
    /// let mut entities = Entities::new(Some(10));
    /// // entity_1 is functionally the entity ID
    /// let entity_1 = entities.get_free_slot().unwrap();
    /// # assert!(entities.add_part(entity_1, Test1 { x: 0 }).is_ok());
    ///
    /// // Because we later need a ref to the same `Type` of map, the mut ref
    /// // will need to be scoped. If the later ref was of a different type,
    /// // eg: Vector2, then it wouldn't need scoping.
    /// {
    ///     let mut part_map = entities.get_pmap_mut::<Test1>().unwrap().unwrap();
    ///     for id in 0..5 {
    ///         if let Ok(part) = part_map.get_part_mut::<Test1>(id) {
    ///             part.x = 42;
    ///         }
    ///     }
    /// }
    ///
    /// // Now get a ref to the modified part
    /// let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
    /// let part = partmap.get_part_ref::<Test1>(entity_1).unwrap();
    /// assert_eq!(part.x, 42);
    /// ```
    pub fn get_pmap_mut<T>(&self) -> Option<Result<RefMut<PartMap>, BorrowMutError>>
    where
        T: 'static,
    {
        if let Some(type_mask) = self.type_masks.get(&TypeId::of::<T>()) {
            return Some(self.parts.get(&type_mask).unwrap().try_borrow_mut());
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use crate::Entities;

    #[test]
    fn add_entity_and_part_check_contains_and_mask() {
        struct Test1 {}

        let mut entities = Entities::new(Some(10));
        // Creating a new entity should find the next free slot
        // and then insert a new blank bitmask
        // Return entity number
        let entity_1 = entities.get_free_slot().unwrap();
        assert_eq!(entity_1, 0);

        assert!(entities.add_part(entity_1, Test1 {}).is_ok());

        let entity_1_mask = entities.get_entity_mask(entity_1);
        // The mask for this entity should have been updated
        let type_mask = entities.get_type_mask::<Test1>().unwrap();
        assert_eq!(entity_1_mask, type_mask);
        assert!(entities.entity_contains::<Test1>(entity_1));

        // should increment entities.next_free_entity
        let e2 = entities.get_free_slot().unwrap();
        assert_eq!(e2, 1);
    }

    #[test]
    fn get_part_map_for_type_as_ref() {
        struct Test1 {}

        let mut entities = Entities::new(Some(10));
        // Should return None
        {
            let m = entities.get_pmap_ref::<Test1>();
            assert_eq!(m.is_none(), true);
        }

        let entity_1 = entities.get_free_slot().unwrap();
        assert!(entities.add_part(entity_1, Test1 {}).is_ok());
        // Should return the HashMap
        let map = entities.get_pmap_ref::<Test1>();
        assert_eq!(map.is_some(), true);
        let map = map.unwrap().unwrap();
        assert!(map.get_part_ref::<Test1>(entity_1).is_ok());
    }

    #[test]
    fn get_part_map_for_type_as_mut() {
        struct Test1 {}

        let mut entities = Entities::new(Some(10));
        // Should return None
        {
            let m = entities.get_pmap_mut::<Test1>();
            assert_eq!(m.is_none(), true);
        }

        let entity_1 = entities.get_free_slot().unwrap();
        assert!(entities.add_part(entity_1, Test1 {}).is_ok());
        // Should return the HashMap
        let map = entities.get_pmap_mut::<Test1>();
        assert_eq!(map.is_some(), true);
        let mut map = map.unwrap().unwrap();
        assert!(map.get_part_mut::<Test1>(entity_1).is_ok());
    }

    #[test]
    fn partmap_ref_and_mut_diff_parts_and_modify() {
        #[derive(Debug, PartialEq)]
        struct Test1 {
            x: u32
        }

        #[derive(Debug, PartialEq)]
        struct Test2 {
            x: u32,
        }

        let mut entities = Entities::new(Some(10));
        let entity_1 = entities.get_free_slot().unwrap();
        assert!(entities.add_part(entity_1, Test1 { x: 66 }).is_ok());
        assert!(entities.add_part(entity_1, Test2 { x: 42 }).is_ok());

        let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
        let part = partmap.get_part_ref::<Test1>(entity_1).unwrap();
        assert_eq!(part, &Test1 { x: 66 });
        assert!(entities.entity_contains::<Test1>(entity_1));

        let mut partmap = entities.get_pmap_mut::<Test2>().unwrap().unwrap();
        let part = partmap.get_part_mut::<Test2>(entity_1).unwrap();
        assert_eq!(part.x, 42);
        part.x = 666;
        assert_ne!(part.x, 42);
    }

    #[test]
    fn partmap_iter_ref() {
        struct Test1 { x: u32 }

        let mut entities = Entities::new(Some(10));
        let mut ids = Vec::new();
        // 1
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);
        // 2
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);
        // 3
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);

        let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
        for (k, v) in partmap.iter_ref() {
            let part = v.downcast_ref::<Test1>().unwrap();
            assert!(part.x == *k);
        }
    }

    #[test]
    fn partmap_iter_mut() {
        struct Test1 { x: u32 }

        let mut entities = Entities::new(Some(10));
        let mut ids = Vec::new();
        // 1
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);
        // 2
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);
        // 3
        let id = entities.get_free_slot().unwrap();
        assert!(entities.add_part(id, Test1 { x: id }).is_ok());
        ids.push(id);

        let mut partmap = entities.get_pmap_mut::<Test1>().unwrap().unwrap();
        for (k, v) in partmap.iter_mut() {
            let part = v.downcast_mut::<Test1>().unwrap();
            part.x += 1;
            assert!(part.x > *k);
        }
    }

    #[test]
    fn remove_parts_then_entity() {
        #[derive(Debug, PartialEq)]
        struct Test1 {}

        #[derive(Debug, PartialEq)]
        struct Test2 {
            x: u32,
        }

        let mut entities = Entities::new(Some(10));
        let entity_1 = entities.get_free_slot().unwrap();
        assert!(entities.add_part(entity_1, Test1 {}).is_ok());
        assert!(entities.add_part(entity_1, Test2 { x: 42 }).is_ok());

        // rm and check entity components
        assert!(entities.entity_contains::<Test1>(entity_1));
        assert!(entities.rm_part::<Test1>(entity_1));
        assert!(!entities.entity_contains::<Test1>(entity_1));
        assert!(entities.entity_contains::<Test2>(entity_1));

        // check masks
        let type_mask_2 = entities.get_type_mask::<Test2>().unwrap();
        assert_eq!(entities.get_entity_mask(entity_1), type_mask_2);

        // Removing all parts erases the entity
        assert!(entities.rm_part::<Test2>(entity_1));
        assert_eq!(entities.get_entity_mask(entity_1), 0);
    }

    #[test]
    fn removing_entity() {
        struct Test1 {}

        let mut entities = Entities::new(Some(10));
        let entity_1 = entities.get_free_slot().unwrap();
        assert!(entities.add_part(entity_1, Test1 {}).is_ok());

        entities.rm_entity(entity_1);

        let partmap = entities.get_pmap_ref::<Test1>().unwrap().unwrap();
        assert_eq!(partmap.get_part_ref::<Test1>(entity_1).is_err(), true);
    }
}