path: root/ecs/src/relationship.rs

use std::any::type_name;
use std::marker::PhantomData;

use ecs_macros::Component;

use crate::component::storage::Storage as ComponentStorage;
use crate::component::{
    Component,
    FromOptional as FromOptionalComponent,
    FromOptionalMut as FromOptionalMutComponent,
};
use crate::lock::ReadGuard;
use crate::system::{ComponentRef, ComponentRefMut};
use crate::uid::{Kind as UidKind, Uid};
use crate::World;

/// A relationship to one or more targets.
#[derive(Debug, Component)]
#[component(
    ref_type = Relation<'component, Kind, ComponentT>,
    ref_mut_type = RelationMut<'component, Kind, ComponentT>,
)]
pub struct Relationship<Kind, ComponentT: Component>
where
    Kind: 'static,
{
    entity_uid: SingleOrMultiple<Uid>,
    _pd: PhantomData<(Kind, ComponentT)>,
}

impl<Kind, ComponentT> Relationship<Kind, ComponentT>
where
    ComponentT: Component,
{
    /// Creates a new `Relationship` with a single target.
    #[must_use]
    pub fn new(entity_uid: Uid) -> Self
    {
        debug_assert_eq!(entity_uid.kind(), UidKind::Entity);

        Self {
            entity_uid: SingleOrMultiple::Single(entity_uid),
            _pd: PhantomData,
        }
    }

    /// Creates a new `Relationship` with multiple targets.
    #[must_use]
    pub fn new_multiple(entity_uids: impl IntoIterator<Item = Uid>) -> Self
    {
        let uids = entity_uids.into_iter().collect::<Vec<_>>();

        for euid in &uids {
            debug_assert_eq!(euid.kind(), UidKind::Entity);
        }

        Self {
            entity_uid: SingleOrMultiple::Multiple(uids),
            _pd: PhantomData,
        }
    }
}

pub struct RelationMut<'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    component_storage_lock: ReadGuard<'static, ComponentStorage>,
    relationship_comp: ComponentRefMut<'rel_comp, Relationship<Kind, ComponentT>>,
}

impl<'rel_comp, Kind, ComponentT> FromOptionalMutComponent<'rel_comp>
    for RelationMut<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    fn from_optional_mut_component(
        optional_component: Option<
            crate::lock::WriteGuard<'rel_comp, Box<dyn Component>>,
        >,
        world: &'rel_comp World,
    ) -> Self
    {
        let relationship_comp =
            ComponentRefMut::<Relationship<Kind, ComponentT>>::from_optional_mut_component(
                optional_component,
                world,
            );

        let component_storage_lock = world
            .data
            .component_storage
            .read_nonblock()
            .expect("Failed to aquire read-only component storage lock");

        Self {
            relationship_comp,
            // SAFETY: The component lock is not used for longer than the original
            // lifetime
            component_storage_lock: unsafe { component_storage_lock.upgrade_lifetime() },
        }
    }
}

impl<'rel_comp, Kind, ComponentT> RelationMut<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    /// Returns the component of the target at the specified index.
    ///
    /// # Panics
    /// Will panic if the entity does not exist in the archetype it belongs to. This
    /// should hopefully never happend.
    #[must_use]
    pub fn get(&self, index: usize) -> Option<ComponentRefMut<'_, ComponentT>>
    {
        let target = self.get_target(index)?;

        let archetype = self.component_storage_lock.get_entity_archetype(*target)?;

        let entity = archetype
            .get_entity(*target)
            .expect("Target entity is gone from archetype");

        let component_index = archetype.get_index_for_component(ComponentT::id())?;

        let component = ComponentRefMut::new(
            entity
                .get_component(component_index)?
                .component
                .write_nonblock()
                .unwrap_or_else(|_| {
                    panic!(
                        "Failed to aquire read-write lock of component {}",
                        type_name::<ComponentT>()
                    )
                }),
        );

        Some(component)
    }

    /// Returns a reference to the target at the specified index.
    #[must_use]
    pub fn get_target(&self, index: usize) -> Option<&Uid>
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get(index),
            SingleOrMultiple::Single(_) => None,
        }
    }

    /// Returns a mutable reference to the target at the specified index.
    #[must_use]
    pub fn get_target_mut(&mut self, index: usize) -> Option<&mut Uid>
    {
        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get_mut(index),
            SingleOrMultiple::Single(_) => None,
        }
    }

    /// Adds a target to the relationship.
    pub fn add_target(&mut self, entity_uid: Uid)
    {
        debug_assert_eq!(entity_uid.kind(), UidKind::Entity);

        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(prev_entity_uid) => {
                self.relationship_comp.entity_uid =
                    SingleOrMultiple::Multiple(vec![*prev_entity_uid, entity_uid]);
            }
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.push(entity_uid),
        }
    }

    /// Removes a target to the relationship, returning it.
    pub fn remove_target(&mut self, index: usize) -> Option<Uid>
    {
        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) => {
                let prev_entity_uid = *entity_uid;

                self.relationship_comp.entity_uid =
                    SingleOrMultiple::Multiple(Vec::new());

                Some(prev_entity_uid)
            }
            SingleOrMultiple::Multiple(entity_uids) => {
                if index >= entity_uids.len() {
                    return None;
                }

                Some(entity_uids.remove(index))
            }
        }
    }

    #[must_use]
    pub fn target_count(&self) -> usize
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(_) => 1,
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.len(),
        }
    }

    /// Returns a iterator of the components of the targets of this relationship.
    #[must_use]
    pub fn iter(&self) -> TargetComponentIterMut<'_, 'rel_comp, Kind, ComponentT>
    {
        TargetComponentIterMut { relation: self, index: 0 }
    }
}

impl<'relationship, 'rel_comp, Kind, ComponentT> IntoIterator
    for &'relationship RelationMut<'rel_comp, Kind, ComponentT>
where
    'relationship: 'rel_comp,
    ComponentT: Component,
{
    type IntoIter = TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>;
    type Item = ComponentRefMut<'rel_comp, ComponentT>;

    fn into_iter(self) -> Self::IntoIter
    {
        self.iter()
    }
}

/// Iterator of the components of the targets of a relationship.
pub struct TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    relation: &'relationship RelationMut<'rel_comp, Kind, ComponentT>,
    index: usize,
}

impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
    for TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>
where
    'relationship: 'rel_comp,
    Kind: 'static,
    ComponentT: Component,
{
    type Item = ComponentRefMut<'rel_comp, ComponentT>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let index = self.index;

        self.index += 1;

        self.relation.get(index)
    }
}

#[derive(Debug)]
enum SingleOrMultiple<Value>
{
    Single(Value),
    Multiple(Vec<Value>),
}

pub struct Relation<'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    component_storage_lock: ReadGuard<'static, ComponentStorage>,
    relationship_comp: ComponentRef<'rel_comp, Relationship<Kind, ComponentT>>,
}

impl<'rel_comp, Kind, ComponentT> FromOptionalComponent<'rel_comp>
    for Relation<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    fn from_optional_component(
        optional_component: Option<ReadGuard<'rel_comp, Box<dyn Component>>>,
        world: &'rel_comp World,
    ) -> Self
    {
        let relationship_comp =
            ComponentRef::<Relationship<Kind, ComponentT>>::from_optional_component(
                optional_component,
                world,
            );

        let component_storage_lock = world
            .data
            .component_storage
            .read_nonblock()
            .expect("Failed to aquire read-only component storage lock");

        Self {
            relationship_comp,
            // SAFETY: The component lock is not used for longer than the original
            // lifetime
            component_storage_lock: unsafe { component_storage_lock.upgrade_lifetime() },
        }
    }
}

impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    /// Returns the component of the target at the specified index.
    ///
    /// # Panics
    /// Will panic if the entity does not exist in the archetype it belongs to. This
    /// should hopefully never happend.
    #[must_use]
    pub fn get(&self, index: usize) -> Option<ComponentRef<'_, ComponentT>>
    {
        let target = self.get_target(index)?;

        let archetype = self.component_storage_lock.get_entity_archetype(*target)?;

        let entity = archetype
            .get_entity(*target)
            .expect("Target entity is gone from archetype");

        let component_index = archetype.get_index_for_component(ComponentT::id())?;

        let component = ComponentRef::new(
            entity
                .get_component(component_index)?
                .component
                .read_nonblock()
                .unwrap_or_else(|_| {
                    panic!(
                        "Failed to aquire read-write lock of component {}",
                        type_name::<ComponentT>()
                    )
                }),
        );

        Some(component)
    }

    /// Returns a reference to the target at the specified index.
    #[must_use]
    pub fn get_target(&self, index: usize) -> Option<&Uid>
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get(index),
            SingleOrMultiple::Single(_) => None,
        }
    }

    #[must_use]
    pub fn target_count(&self) -> usize
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(_) => 1,
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.len(),
        }
    }

    pub fn target_uids(&self) -> impl Iterator<Item = Uid> + '_
    {
        (0..self.target_count())
            .map_while(|target_index| self.get_target(target_index).copied())
    }

    /// Returns a iterator of the components of the targets of this relationship.
    #[must_use]
    pub fn iter(&self) -> TargetComponentIter<'_, 'rel_comp, Kind, ComponentT>
    {
        TargetComponentIter { relation: self, index: 0 }
    }
}

impl<'relationship, 'rel_comp, Kind, ComponentT> IntoIterator
    for &'relationship Relation<'rel_comp, Kind, ComponentT>
where
    'relationship: 'rel_comp,
    ComponentT: Component,
{
    type IntoIter = TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>;
    type Item = ComponentRef<'rel_comp, ComponentT>;

    fn into_iter(self) -> Self::IntoIter
    {
        self.iter()
    }
}

/// Iterator of the components of the targets of a relationship.
pub struct TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    relation: &'relationship Relation<'rel_comp, Kind, ComponentT>,
    index: usize,
}

impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
    for TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
where
    'relationship: 'rel_comp,
    Kind: 'static,
    ComponentT: Component,
{
    type Item = ComponentRef<'rel_comp, ComponentT>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let index = self.index;

        self.index += 1;

        self.relation.get(index)
    }
}