path: root/ecs/src/component/storage/graph.rs

use hashbrown::hash_map::Values as HashMapValueIter;
use hashbrown::HashMap;

use crate::component::storage::archetype::{Archetype, Id as ArchetypeId};
use crate::uid::{Kind as UidKind, Uid};

#[derive(Debug, Default)]
pub struct Graph
{
    nodes: Vec<ArchetypeNode>,
    archetype_index_lookup: HashMap<ArchetypeId, usize>,
}

impl Graph
{
    pub fn create_node(&mut self, id: ArchetypeId, component_ids: &impl AsRef<[Uid]>)
    {
        debug_assert!(!self.contains_archetype(id));

        let _ = self.get_or_create_node(id, component_ids);
    }

    pub fn get_or_create_node(
        &mut self,
        id: ArchetypeId,
        component_ids: &impl AsRef<[Uid]>,
    ) -> &mut ArchetypeNode
    {
        let index = *self.archetype_index_lookup.entry(id).or_insert_with(|| {
            self.nodes.push(ArchetypeNode {
                archetype: Archetype::new(id, component_ids.as_ref()),
                edges: HashMap::new(),
            });

            self.nodes.len() - 1
        });

        self.create_missing_edges(id);

        self.nodes
            .get_mut(index)
            .expect("Archetype index from lookup is out of bounds")
    }

    pub fn contains_archetype(&self, id: ArchetypeId) -> bool
    {
        self.archetype_index_lookup.contains_key(&id)
    }

    pub fn get_node_by_id(&self, id: ArchetypeId) -> Option<&ArchetypeNode>
    {
        let index = self.archetype_index_lookup.get(&id)?;

        Some(self.nodes.get(*index).unwrap_or_else(|| {
            panic!("In invalid state! Index of archetype with ID {id:?} is out of bounds")
        }))
    }

    pub fn get_node_by_id_mut(&mut self, id: ArchetypeId) -> Option<&mut ArchetypeNode>
    {
        let index = self.archetype_index_lookup.get(&id)?;

        Some(self.nodes.get_mut(*index).unwrap_or_else(|| {
            panic!("In invalid state! Index of archetype with ID {id:?} is out of bounds")
        }))
    }

    pub fn recurse_archetype_add_edges(
        &self,
        archetype_id: ArchetypeId,
    ) -> Option<ArchetypeAddEdgeRecursiveIter>
    {
        Some(ArchetypeAddEdgeRecursiveIter {
            graph: self,
            stack: vec![self.get_node_by_id(archetype_id)?.edges.values()],
        })
    }

    fn create_missing_edges(&mut self, archetype_id: ArchetypeId)
    {
        let archetype_node = self
            .get_node_by_id(archetype_id)
            .expect("Archetype should exist");

        let mut comp_ids = archetype_node
            .archetype()
            .component_ids_unsorted()
            .collect::<Vec<_>>();

        comp_ids.sort();

        for component_id_index in 0..archetype_node.archetype().component_cnt() {
            let mut comp_ids_without_comp = comp_ids.clone();

            let removed_comp_id = comp_ids_without_comp.remove(component_id_index);

            let archetype_without_comp_id = ArchetypeId::new(&comp_ids_without_comp);

            let Some(archetype_node_without_comp) =
                self.get_node_by_id_mut(archetype_without_comp_id)
            else {
                continue;
            };

            archetype_node_without_comp
                .get_or_insert_edges(removed_comp_id, ArchetypeEdges::default)
                .add = Some(archetype_id);

            let archetype_node = self
                .get_node_by_id_mut(archetype_id)
                .expect("Archetype should exist");

            archetype_node
                .get_or_insert_edges(removed_comp_id, ArchetypeEdges::default)
                .remove = Some(archetype_without_comp_id);
        }

        let archetype_node_index = *self
            .archetype_index_lookup
            .get(&archetype_id)
            .expect("Archetype should exist");

        let (nodes_before, nodes_rest) = self.nodes.split_at_mut(archetype_node_index);

        let ([archetype_node], nodes_after) = nodes_rest.split_at_mut(1) else {
            unreachable!();
        };

        let expected_component_cnt = archetype_node.archetype().component_cnt() + 1;

        for other_archetype_node in nodes_before.iter_mut().chain(nodes_after.iter_mut())
        {
            if other_archetype_node.archetype().component_cnt() != expected_component_cnt
                || !other_archetype_node
                    .archetype()
                    .is_superset(&archetype_node.archetype())
            {
                continue;
            }

            let extra_comp_id = other_archetype_node
                .archetype()
                .component_ids_unsorted()
                .find(|comp_id| {
                    !archetype_node.archetype().has_component_with_id(*comp_id)
                })
                .expect("Archetype should contain one extra component ID");

            other_archetype_node
                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
                .remove = Some(archetype_id);

            archetype_node
                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
                .add = Some(other_archetype_node.archetype().id());
        }
    }
}

#[derive(Debug)]
pub struct ArchetypeNode
{
    archetype: Archetype,
    edges: HashMap<Uid, ArchetypeEdges>,
}

impl ArchetypeNode
{
    pub fn archetype(&self) -> &Archetype
    {
        &self.archetype
    }

    pub fn archetype_mut(&mut self) -> &mut Archetype
    {
        &mut self.archetype
    }

    pub fn get_or_insert_edges(
        &mut self,
        component_id: Uid,
        insert_fn: impl FnOnce() -> ArchetypeEdges,
    ) -> &mut ArchetypeEdges
    {
        debug_assert_eq!(component_id.kind(), UidKind::Component);

        self.edges.entry(component_id).or_insert_with(insert_fn)
    }

    pub fn get_edges_mut(&mut self, component_id: Uid) -> Option<&mut ArchetypeEdges>
    {
        debug_assert_eq!(component_id.kind(), UidKind::Component);

        self.edges.get_mut(&component_id)
    }

    pub fn make_add_edge(&self, component_id: Uid) -> (ArchetypeId, Vec<Uid>)
    {
        let mut edge_comp_ids = self
            .archetype()
            .component_ids_unsorted()
            .chain([component_id])
            .collect::<Vec<_>>();

        edge_comp_ids.sort();

        let add_edge_id = ArchetypeId::new(&edge_comp_ids);

        (add_edge_id, edge_comp_ids)
    }

    pub fn make_remove_edge(&self, component_id: Uid) -> (ArchetypeId, Vec<Uid>)
    {
        let mut edge_comp_ids = self
            .archetype()
            .component_ids_unsorted()
            .filter(|id| *id != component_id)
            .collect::<Vec<_>>();

        edge_comp_ids.sort();

        let remove_edge_id = ArchetypeId::new(&edge_comp_ids);

        (remove_edge_id, edge_comp_ids)
    }
}

#[derive(Debug, Default)]
pub struct ArchetypeEdges
{
    pub add: Option<ArchetypeId>,
    pub remove: Option<ArchetypeId>,
}

#[derive(Debug)]
pub struct ArchetypeAddEdgeRecursiveIter<'graph>
{
    graph: &'graph Graph,
    stack: Vec<HashMapValueIter<'graph, Uid, ArchetypeEdges>>,
}

impl<'graph> Iterator for ArchetypeAddEdgeRecursiveIter<'graph>
{
    type Item = Option<ArchetypeId>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let iter = self.stack.last_mut()?;

        let Some(edges) = iter.next() else {
            self.stack.pop();

            return Some(None);
        };

        let Some(add_edge) = edges.add else {
            return Some(None);
        };

        let add_edge_archetype = self
            .graph
            .get_node_by_id(add_edge)
            .expect("Add edge archetype not found");

        self.stack.push(add_edge_archetype.edges.values());

        Some(Some(add_edge))
    }
}