use std::any::type_name;
use std::array::IntoIter as ArrayIter;
use std::cell::RefCell;
use std::vec::IntoIter as VecIntoIter;

use hashbrown::HashMap;

use crate::component::storage::archetype::{
    Archetype,
    ArchetypeEntity,
    Id as ArchetypeId,
};
use crate::component::storage::graph::{
    ArchetypeAddEdgeDfsIter,
    ArchetypeAddEdgeDfsIterResult,
    ArchetypeEdges,
    Graph,
};
use crate::component::{Component, Metadata as ComponentMetadata};
use crate::type_name::TypeName;
use crate::uid::{Kind as UidKind, Uid};
use crate::util::{BorrowedOrOwned, Either, StreamingIterator, VecExt};
use crate::EntityComponent;

pub mod archetype;

mod graph;

#[derive(Debug, Default)]
pub struct Storage
{
    graph: Graph,
    entity_archetype_lookup: HashMap<Uid, ArchetypeId>,
    imaginary_archetypes: RefCell<Vec<ImaginaryArchetype>>,
}

impl Storage
{
    pub fn search_archetypes(
        &self,
        component_metadata: &[ComponentMetadata],
    ) -> ArchetypeRefIter<'_>
    {
        let archetype_id = ArchetypeId::from_components_metadata(component_metadata);

        let Some(add_edge_recursive_iter) =
            self.graph.dfs_archetype_add_edges(archetype_id)
        else {
            self.imaginary_archetypes
                .borrow_mut()
                .push(ImaginaryArchetype {
                    id: archetype_id,
                    component_ids: component_metadata
                        .iter()
                        .filter_map(|comp_metadata| {
                            if comp_metadata.is_optional {
                                return None;
                            }

                            Some(comp_metadata.id)
                        })
                        .collect::<Vec<_>>(),
                });

            let found_archetypes = self.find_all_archetype_with_comps(component_metadata);

            return ArchetypeRefIter {
                storage: self,
                pre_iter: Either::B(found_archetypes.clone().into_iter()),
                dfs_iter: ArchetypeAddEdgeDfsIter::new(&self.graph, &found_archetypes),
            };
        };

        ArchetypeRefIter {
            storage: self,
            pre_iter: Either::A([archetype_id].into_iter()),
            dfs_iter: add_edge_recursive_iter,
        }
    }

    pub fn get_archetype_by_id(&self, id: ArchetypeId) -> Option<&Archetype>
    {
        Some(self.graph.get_node_by_id(id)?.archetype())
    }

    pub fn create_entity(&mut self, uid: Uid) -> Result<(), Error>
    {
        debug_assert_eq!(uid.kind(), UidKind::Entity);

        if self.entity_archetype_lookup.contains_key(&uid) {
            return Err(Error::EntityAlreadyExists(uid));
        }

        let empty_archetype_id = ArchetypeId::from_components_metadata(&[]);

        let archetype_node = self.graph.get_or_create_node(empty_archetype_id, &[]);

        archetype_node
            .archetype_mut()
            .push_entity(ArchetypeEntity { uid, components: vec![] });

        self.entity_archetype_lookup.insert(uid, empty_archetype_id);

        Ok(())
    }

    pub fn remove_entity(
        &mut self,
        entity_uid: Uid,
    ) -> Result<Vec<EntityComponent>, Error>
    {
        let Some(archetype_id) = self.entity_archetype_lookup.get(&entity_uid) else {
            return Err(Error::EntityDoesNotExist(entity_uid));
        };

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

        let entity = archetype_node
            .archetype_mut()
            .remove_entity(entity_uid)
            .expect("Entity should exist in archetype");

        self.entity_archetype_lookup.remove(&entity_uid);

        Ok(entity.components)
    }

    pub fn get_entity_archetype(&self, entity_uid: Uid) -> Option<&Archetype>
    {
        let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;

        self.get_archetype_by_id(*archetype_id)
    }

    pub fn add_entity_component(
        &mut self,
        entity_uid: Uid,
        component: Box<dyn Component>,
    ) -> Result<(), Error>
    {
        debug_assert!(
            !component.self_is_optional(),
            "Adding a optional component to a entity is not supported"
        );

        let Some(archetype_id) = self.entity_archetype_lookup.get(&entity_uid) else {
            return Err(Error::EntityDoesNotExist(entity_uid));
        };

        let archetype_id = *archetype_id;

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

        if archetype_node
            .archetype()
            .has_component_with_id(component.self_id())
        {
            return Err(Error::ComponentAlreadyInEntity {
                entity: entity_uid,
                component: component.self_id(),
            });
        }

        let add_edge_archetype_id = archetype_node
            .get_or_insert_edges(component.self_id(), ArchetypeEdges::default)
            .add
            .unwrap_or_else(|| {
                let archetype_node = self
                    .graph
                    .get_node_by_id_mut(archetype_id)
                    .expect("Archetype should exist");

                let (add_edge_id, add_edge_comp_ids) =
                    archetype_node.make_add_edge(component.self_id());

                archetype_node
                    .get_edges_mut(component.self_id())
                    .expect("Edges for component in archetype should exist")
                    .add = Some(add_edge_id);

                if !self.graph.contains_archetype(add_edge_id) {
                    self.graph.create_node(add_edge_id, &add_edge_comp_ids);
                }

                add_edge_id
            });

        {
            let add_edge_archetype_node = self
                .graph
                .get_node_by_id_mut(add_edge_archetype_id)
                .expect("Add edge archetype should exist");

            let add_edge_archetype_edges = add_edge_archetype_node
                .get_or_insert_edges(component.self_id(), ArchetypeEdges::default);

            add_edge_archetype_edges.remove = Some(archetype_id);
        }

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

        let mut entity = archetype_node
            .archetype_mut()
            .remove_entity(entity_uid)
            .expect("Entity should exist in archetype");

        let add_edge_archetype = self
            .graph
            .get_node_by_id_mut(add_edge_archetype_id)
            .expect("Add edge archetype should exist")
            .archetype_mut();

        let component_index = add_edge_archetype
            .get_index_for_component(component.self_id())
            .expect("Archetype should have index for component");

        entity.components.insert(component_index, component.into());

        add_edge_archetype.push_entity(entity);

        self.entity_archetype_lookup
            .insert(entity_uid, add_edge_archetype_id);

        Ok(())
    }

    pub fn remove_entity_component(
        &mut self,
        entity_uid: Uid,
        component_id: Uid,
    ) -> Result<(), Error>
    {
        let Some(archetype_id) = self.entity_archetype_lookup.get(&entity_uid) else {
            return Err(Error::EntityDoesNotExist(entity_uid));
        };

        let archetype_id = *archetype_id;

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

        if !archetype_node
            .archetype()
            .has_component_with_id(component_id)
        {
            return Err(Error::ComponentNotFoundInEntity {
                entity: entity_uid,
                component: component_id,
            });
        }

        let remove_edge_id = archetype_node
            .get_or_insert_edges(component_id, ArchetypeEdges::default)
            .remove
            .unwrap_or_else(|| {
                let archetype_node = self
                    .graph
                    .get_node_by_id_mut(archetype_id)
                    .expect("Archetype should exist");

                let (remove_edge_id, remove_edge_comp_ids) =
                    archetype_node.make_remove_edge(component_id);

                archetype_node
                    .get_edges_mut(component_id)
                    .expect("Edges for component in archetype should exist")
                    .remove = Some(remove_edge_id);

                if !self.graph.contains_archetype(remove_edge_id) {
                    self.graph
                        .create_node(remove_edge_id, &remove_edge_comp_ids);
                }

                remove_edge_id
            });

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

        let mut entity = archetype_node
            .archetype_mut()
            .remove_entity(entity_uid)
            .expect("Entity should exist in archetype");

        let (comp_to_remove_index, _) = entity
            .components
            .iter()
            .enumerate()
            .find(|(_, comp)| comp.id == component_id)
            .expect("Entity should contain component");

        entity.components.remove(comp_to_remove_index);

        self.graph
            .get_node_by_id_mut(remove_edge_id)
            .expect("Remove edge archetype should exist")
            .archetype_mut()
            .push_entity(entity);

        self.entity_archetype_lookup
            .insert(entity_uid, remove_edge_id);

        Ok(())
    }

    pub fn create_imaginary_archetypes(&mut self)
    {
        for imaginary_archetype in self.imaginary_archetypes.get_mut().drain(..) {
            if self.graph.contains_archetype(imaginary_archetype.id) {
                continue;
            }

            self.graph
                .create_node(imaginary_archetype.id, &imaginary_archetype.component_ids);
        }
    }

    fn find_all_archetype_with_comps(
        &self,
        component_metadata: &[ComponentMetadata],
    ) -> Vec<ArchetypeId>
    {
        let comp_metadata_not_opt_cnt = component_metadata
            .iter()
            .filter(|comp_metadata| !comp_metadata.is_optional)
            .count();

        let Some(mut search_iter) =
            self.graph.dfs_archetype_add_edges(ArchetypeId::new(&[]))
        else {
            // If the root archetype doesn't exist, no other archetype can exist either
            return Vec::new();
        };

        let mut found = Vec::<ArchetypeId>::new();

        while let Some(node_id) = search_iter.streaming_next() {
            let ArchetypeAddEdgeDfsIterResult::AddEdge(node_id) = node_id else {
                continue;
            };

            let node = self
                .graph
                .get_node_by_id(node_id)
                .expect("Graph node found through DFS doesn't exist");

            if node.archetype().component_cnt() < comp_metadata_not_opt_cnt {
                continue;
            }

            if !component_metadata
                .iter()
                .filter(|comp_metadata| !comp_metadata.is_optional)
                .all(|comp_metadata| {
                    node.archetype().has_component_with_id(comp_metadata.id)
                })
            {
                continue;
            }

            found.push(node.archetype().id());

            search_iter.pop();
        }

        found
    }
}

impl TypeName for Storage
{
    fn type_name(&self) -> &'static str
    {
        type_name::<Self>()
    }
}

#[derive(Debug)]
pub struct ArchetypeRefIter<'storage>
{
    storage: &'storage Storage,
    pre_iter: Either<ArrayIter<ArchetypeId, 1>, VecIntoIter<ArchetypeId>>,
    dfs_iter: ArchetypeAddEdgeDfsIter<'storage>,
}

impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
{
    type Item = &'component_storage Archetype;

    fn next(&mut self) -> Option<Self::Item>
    {
        if let Some(pre_iter_archetype_id) = self.pre_iter.next() {
            return Some(
                self.storage
                    .get_archetype_by_id(pre_iter_archetype_id)
                    .expect("Archetype should exist"),
            );
        }

        let archetype_id = loop {
            match self.dfs_iter.streaming_find(|res| {
                matches!(
                    res,
                    ArchetypeAddEdgeDfsIterResult::AddEdge { .. }
                        | ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound { .. }
                )
            })? {
                ArchetypeAddEdgeDfsIterResult::AddEdge(add_edge_archetype_id) => {
                    break add_edge_archetype_id
                }
                ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound {
                    archetype,
                    add_edge_archetype_id,
                    add_edge_component_id,
                } => {
                    let mut add_edge_archetype_comps = archetype
                        .component_ids_sorted()
                        .map(|id| ComponentMetadata { id, is_optional: false })
                        .collect::<Vec<_>>();

                    add_edge_archetype_comps.insert_at_partition_point_by_key(
                        ComponentMetadata::new_non_optional(add_edge_component_id),
                        |comp_metadata| comp_metadata.id,
                    );

                    self.storage.imaginary_archetypes.borrow_mut().push(
                        ImaginaryArchetype {
                            id: add_edge_archetype_id,
                            component_ids: add_edge_archetype_comps
                                .iter()
                                .map(|comp_metadata| comp_metadata.id)
                                .collect::<Vec<_>>(),
                        },
                    );

                    let found =
                        self.find_edges_of_imaginary_archetype(&add_edge_archetype_comps);

                    self.dfs_iter.push((
                        BorrowedOrOwned::Owned(Archetype::new(
                            add_edge_archetype_id,
                            add_edge_archetype_comps
                                .iter()
                                .map(|metadata| metadata.id)
                                .collect::<Vec<_>>(),
                        )),
                        found.into_iter(),
                    ));

                    continue;
                }
                _ => {
                    unreachable!();
                }
            }
        };

        Some(
            self.storage
                .get_archetype_by_id(archetype_id)
                .expect("Archetype should exist"),
        )
    }
}

impl ArchetypeRefIter<'_>
{
    fn find_edges_of_imaginary_archetype(
        &self,
        imaginary_archetype_comps: &[ComponentMetadata],
    ) -> Vec<(Uid, ArchetypeEdges)>
    {
        self.storage
            .find_all_archetype_with_comps(imaginary_archetype_comps)
            .into_iter()
            .filter_map(|found_id| {
                let found_archetype = self.storage.get_archetype_by_id(found_id).unwrap();

                if found_archetype.component_cnt() < imaginary_archetype_comps.len() + 1 {
                    return None;
                }

                let unique_comp_id = found_archetype
                    .component_ids_sorted()
                    .find(|id| {
                        !imaginary_archetype_comps
                            .iter()
                            .any(|comp_metadata| comp_metadata.id == *id)
                    })
                    .expect("Oh noooo");

                let mut add_edge_comp_ids = imaginary_archetype_comps
                    .iter()
                    .map(|comp_metadata| comp_metadata.id)
                    .collect::<Vec<_>>();

                add_edge_comp_ids
                    .insert_at_partition_point_by_key(unique_comp_id, |id| *id);

                let add_edge = ArchetypeId::new(&add_edge_comp_ids);

                Some((
                    unique_comp_id,
                    ArchetypeEdges { add: Some(add_edge), remove: None },
                ))
            })
            .collect::<Vec<_>>()
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error
{
    #[error("Entity with ID {0:?} already exists")]
    EntityAlreadyExists(Uid),

    #[error("Entity with ID {0:?} does not exist")]
    EntityDoesNotExist(Uid),

    #[error("Entity with ID {entity:?} already has component with ID {component:?}")]
    ComponentAlreadyInEntity
    {
        entity: Uid, component: Uid
    },

    #[error("Entity with ID {entity:?} does not have component with ID {component:?}")]
    ComponentNotFoundInEntity
    {
        entity: Uid, component: Uid
    },
}

#[derive(Debug)]
struct ImaginaryArchetype
{
    id: ArchetypeId,
    component_ids: Vec<Uid>,
}

#[cfg(test)]
mod tests
{
    use crate::component::storage::archetype::Id as ArchetypeId;
    use crate::component::storage::Storage;
    use crate::uid::{Kind as UidKind, Uid};

    #[test]
    fn create_entity_works()
    {
        let mut new_storage = Storage::default();

        let uid = Uid::new_unique(UidKind::Entity);

        new_storage.create_entity(uid).expect("Expected Ok");

        let archetype_node = new_storage
            .graph
            .get_node_by_id(ArchetypeId::from_components_metadata(&[]))
            .expect("Archetype for entities with no component doesn't exist");

        assert_eq!(archetype_node.archetype().component_cnt(), 0);
        assert_eq!(archetype_node.archetype().entity_cnt(), 1);

        assert_eq!(
            new_storage.entity_archetype_lookup.get(&uid).copied(),
            Some(ArchetypeId::from_components_metadata(&[]))
        );
    }
}