use std::vec::IntoIter as VecIntoIter;

use hashbrown::{HashMap, HashSet};

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

#[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 exists_before = self.archetype_index_lookup.contains_key(&id);

        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
        });

        if !exists_before {
            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 dfs_archetype_add_edges(
        &self,
        archetype_id: ArchetypeId,
    ) -> Option<ArchetypeAddEdgeDfsIter>
    {
        let node = self.get_node_by_id(archetype_id)?;

        Some(ArchetypeAddEdgeDfsIter {
            graph: self,
            stack: vec![(
                BorrowedOrOwned::Borrowned(node.archetype()),
                node.edges
                    .iter()
                    .map(|(comp_id, edges)| (*comp_id, edges.clone()))
                    .collect::<Vec<_>>()
                    .into_iter(),
            )],
            visited: HashSet::new(),
        })
    }

    fn create_missing_edges(&mut self, archetype_id: ArchetypeId)
    {
        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!();
        };

        for other_archetype_node in nodes_before.iter_mut().chain(nodes_after.iter_mut())
        {
            if archetype_node.archetype().component_cnt()
                > other_archetype_node.archetype().component_cnt()
                && other_archetype_node
                    .archetype()
                    .is_subset(archetype_node.archetype())
            {
                Self::create_missing_subset_node_edges(
                    archetype_node,
                    other_archetype_node,
                );

                continue;
            }

            if other_archetype_node
                .archetype()
                .is_superset(archetype_node.archetype())
            {
                Self::create_missing_superset_node_edges(
                    archetype_node,
                    other_archetype_node,
                );
            }
        }
    }

    fn create_missing_subset_node_edges(
        target_node: &ArchetypeNode,
        subset_node: &mut ArchetypeNode,
    )
    {
        let uniq_comp_id = target_node
            .archetype()
            .component_ids_sorted()
            .find(|id| !subset_node.archetype().has_component_with_id(*id))
            .unwrap();

        subset_node
            .get_or_insert_edges(uniq_comp_id, ArchetypeEdges::default)
            .add = Some(subset_node.make_add_edge(uniq_comp_id).0);
    }

    fn create_missing_superset_node_edges(
        target_node: &mut ArchetypeNode,
        superset_node: &mut ArchetypeNode,
    )
    {
        if superset_node.archetype().component_cnt()
            > target_node.archetype().component_cnt()
        {
            let first_unique_comp_id = superset_node
                .archetype()
                .component_ids_sorted()
                .find(|other_archetype_comp_id| {
                    !target_node
                        .archetype()
                        .has_component_with_id(*other_archetype_comp_id)
                })
                .or_else(|| {
                    if target_node.archetype().component_cnt() != 0 {
                        return None;
                    }

                    superset_node.archetype().component_ids_sorted().next()
                })
                .expect("Not possible");

            target_node
                .get_or_insert_edges(first_unique_comp_id, ArchetypeEdges::default)
                .add = Some(target_node.make_add_edge(first_unique_comp_id).0);

            return;
        }

        if superset_node.archetype().component_cnt()
            != target_node.archetype().component_cnt() + 1
        {
            return;
        }

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

        superset_node
            .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
            .remove = Some(target_node.archetype().id());

        target_node
            .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
            .add = Some(superset_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, Clone)]
pub struct ArchetypeEdges
{
    pub add: Option<ArchetypeId>,
    pub remove: Option<ArchetypeId>,
}

type ArchetypeAddEdgeDfsIterStackElem<'graph> = (
    BorrowedOrOwned<'graph, Archetype>,
    VecIntoIter<(Uid, ArchetypeEdges)>,
);

#[derive(Debug)]
pub struct ArchetypeAddEdgeDfsIter<'graph>
{
    graph: &'graph Graph,
    stack: Vec<ArchetypeAddEdgeDfsIterStackElem<'graph>>,
    visited: HashSet<ArchetypeId>,
}

impl<'graph> ArchetypeAddEdgeDfsIter<'graph>
{
    pub fn new(graph: &'graph Graph, start_nodes: &[ArchetypeId]) -> Self
    {
        Self {
            graph,
            stack: start_nodes
                .iter()
                .map(|start_node_id| {
                    let start_node = graph
                        .get_node_by_id(*start_node_id)
                        .expect("Start node does not exist");

                    (
                        BorrowedOrOwned::Borrowned(start_node.archetype()),
                        start_node
                            .edges
                            .iter()
                            .map(|(comp_id, edges)| (*comp_id, edges.clone()))
                            .collect::<Vec<_>>()
                            .into_iter(),
                    )
                })
                .collect(),
            visited: start_nodes.iter().copied().collect::<HashSet<_>>(),
        }
    }

    pub fn push(
        &mut self,
        item: (
            BorrowedOrOwned<'graph, Archetype>,
            VecIntoIter<(Uid, ArchetypeEdges)>,
        ),
    )
    {
        self.stack.push(item);
    }

    pub fn pop(&mut self)
    {
        self.stack.pop();
    }
}

impl<'graph> StreamingIterator for ArchetypeAddEdgeDfsIter<'graph>
{
    type Item<'a>
        = ArchetypeAddEdgeDfsIterResult<'graph, 'a>
    where
        Self: 'a;

    fn streaming_next(&mut self) -> Option<Self::Item<'_>>
    {
        let (_, edges_iter) = self.stack.last_mut()?;

        let Some((component_id, edges)) = edges_iter.next() else {
            self.stack.pop();

            return Some(ArchetypeAddEdgeDfsIterResult::NoEdgesLeftForArchetype);
        };

        let Some(add_edge) = edges.add else {
            return Some(ArchetypeAddEdgeDfsIterResult::NoAddEdge);
        };

        if self.visited.contains(&add_edge) {
            return Some(ArchetypeAddEdgeDfsIterResult::AddEdgeAlreadyVisited);
        }

        self.visited.insert(add_edge);

        let Some(add_edge_archetype) = self.graph.get_node_by_id(add_edge) else {
            return Some(ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound {
                archetype: &self.stack.last().unwrap().0,
                add_edge_archetype_id: add_edge,
                add_edge_component_id: component_id,
            });
        };

        self.stack.push((
            BorrowedOrOwned::Borrowned(add_edge_archetype.archetype()),
            add_edge_archetype
                .edges
                .iter()
                .map(|(comp_id, edges)| (*comp_id, edges.clone()))
                .collect::<Vec<_>>()
                .into_iter(),
        ));

        Some(ArchetypeAddEdgeDfsIterResult::AddEdge(add_edge))
    }
}

#[derive(Debug)]
pub enum ArchetypeAddEdgeDfsIterResult<'graph, 'iter>
{
    AddEdge(ArchetypeId),
    NoEdgesLeftForArchetype,
    NoAddEdge,
    AddEdgeAlreadyVisited,
    AddEdgeArchetypeNotFound
    {
        archetype: &'iter BorrowedOrOwned<'graph, Archetype>,
        add_edge_archetype_id: ArchetypeId,
        add_edge_component_id: Uid,
    },
}