#[cfg(feature = "serde")]
use crate::serde_utils::*;
use crate::{Aggregate, GraphError, Meta, Metadata};
use anyhow::{anyhow, Error, Result};
use std::collections::{BTreeMap, BTreeSet};
use uuid::Uuid;

/// A node in a graph.
#[derive(Clone, Debug, Default, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(rename_all = "camelCase")
)]
pub struct Node {
    /// Instance metadata.
    pub metadata: Metadata,
    /// ID of the parent of this node, if any.
    #[cfg_attr(
        feature = "serde",
        serde(default, skip_serializing_if = "Option::is_none")
    )]
    pub parent: Option<Uuid>,
    /// IDs of this node's children,
    #[cfg_attr(
        feature = "serde",
        serde(default, skip_serializing_if = "is_empty_vec")
    )]
    pub children: Vec<Uuid>,
    /// Whether this node is a bus node for its parent.
    #[cfg_attr(feature = "serde", serde(default, skip_serializing_if = "is_default"))]
    pub is_bus: bool,
}
impl Node {
    /// Create a new node struct with arguments.
    pub fn new(
        metadata: Option<Metadata>,
        parent: Option<Uuid>,
        children: Option<Vec<Uuid>>,
        is_bus: Option<bool>,
    ) -> Self {
        Node {
            metadata: metadata.unwrap_or_default(),
            parent,
            children: match children {
                None => vec![],
                Some(x) => x,
            },
            is_bus: is_bus.unwrap_or(false),
        }
    }

    /// Whether this node has no parent and is thus a root node.
    pub fn is_root(&self) -> bool {
        self.parent.is_none()
    }

    /// Whether this node has no children and is thus a leaf node.
    pub fn is_leaf(&self) -> bool {
        self.children.is_empty()
    }

    /// Number of children.
    pub fn dim(&self) -> usize {
        self.children.len()
    }
}

impl Meta for Node {
    /// Get node metadata.
    fn get_meta(&self) -> &Metadata {
        &self.metadata
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct NodeStoreData {
    nodes: BTreeMap<Uuid, Node>,
    roots: BTreeSet<Uuid>,
    leafs: BTreeSet<Uuid>,
    aggregate: Aggregate,
    depth: Option<usize>,
}
impl HasNodeStore for NodeStoreData {
    /// Get a reference to the node store.
    fn node_store(&self) -> &NodeStoreData {
        self
    }
    /// Get a mutable reference to the node store.
    fn node_store_mut(&mut self) -> &mut NodeStoreData {
        self
    }
}
impl NodeStore for NodeStoreData {}

pub trait HasNodeStore {
    fn node_store(&self) -> &NodeStoreData;
    fn node_store_mut(&mut self) -> &mut NodeStoreData;
}

pub trait NodeStore: HasNodeStore {
    /// The number of nodes in this store.
    fn nodes_len(&self) -> usize {
        self.node_store().nodes.len()
    }

    /// Maximum node depth or height.
    fn max_node_depth(&mut self) -> usize
    where
        Self: Sized,
    {
        if let Some(depth) = self.node_store().depth {
            return depth;
        }
        let depth = self.calculate_node_depth();
        self.node_store_mut().depth = Some(depth);
        depth
    }

    /// Whether this node store is empty.
    fn is_nodes_empty(&self) -> bool {
        self.nodes_len() == 0
    }

    /// Check whether this store contains a node with the given node ID.
    fn has_node(&self, id: &Uuid) -> bool {
        self.node_store().nodes.contains_key(id)
    }

    /// Add a single node to this store. Setting 'safe' checks whether this node's
    /// references exist and detects cyclic hierarchies.
    fn add_node(&mut self, node: Node, safe: bool) -> Result<Option<Node>>
    where
        Self: Sized,
    {
        if safe {
            self.check_node(&node)?;
        }
        let id = node.id().to_owned();
        let replaced = self.del_node(&id);
        self.node_store_mut().aggregate.add(node.get_meta());
        if node.is_leaf() {
            self.node_store_mut().leafs.insert(id);
        }
        if node.is_root() {
            self.node_store_mut().roots.insert(id);
        }
        if node.parent.is_some() || !node.children.is_empty() {
            self.node_store_mut().depth = None;
        }
        self.node_store_mut().nodes.insert(id, node);
        Ok(replaced)
    }

    /// Extend this node store with new nodes. Setting 'safe' checks whether all node
    /// references exist and detects cyclic hierarchies after adding them to the store
    /// first.
    fn extend_nodes(&mut self, nodes: Vec<Node>, safe: bool) -> Result<Vec<Node>>
    where
        Self: Sized,
    {
        let mut replaced = vec![];
        for node in nodes.clone().into_iter() {
            if let Some(node) = self.add_node(node, false)? {
                replaced.push(node);
            }
        }
        if safe {
            for node in nodes.iter() {
                self.check_node(node)?;
            }
        }
        Ok(replaced)
    }

    /// Check whether a node's references exist and it doesn't partake in cyclic
    /// hierarchies.
    fn check_node(&self, node: &Node) -> Result<()>
    where
        Self: Sized,
    {
        if let Some(parent_id) = &node.parent {
            if !self.has_node(parent_id) {
                return Err(anyhow!(GraphError::NodeNotInStore(parent_id.to_owned())));
            }
        }
        for child_id in node.children.iter() {
            if !self.has_node(child_id) {
                return Err(anyhow!(GraphError::NodeNotInStore(child_id.to_owned())));
            }
        }
        if node.is_bus && node.parent.is_none() {
            return Err(anyhow!(GraphError::BusWithoutParent(node.id().to_owned())));
        }
        // Use query functions in safe mode to detect any cyclic hierarchies.
        // Pretty brute-force, but works.
        self.ascendant_nodes(node.id(), true, true, None, None)
            .try_fold((), |_, r| r.map(|_| ()))?;
        self.descendant_nodes(node.id(), true, true, None, None)
            .try_fold((), |_, r| r.map(|_| ()))?;
        Ok(())
    }

    /// Get a reference to a node.
    fn get_node(&self, id: &Uuid) -> Option<&Node> {
        self.node_store().nodes.get(id)
    }
    /// Get a reference to a node as as result.
    fn get_node_err(&self, id: &Uuid) -> Result<&Node> {
        match self.node_store().nodes.get(id) {
            Some(node) => Ok(node),
            None => Err(anyhow!(GraphError::NodeNotInStore(id.to_owned()))),
        }
    }

    /// Get a mutable reference to a node.
    fn get_node_mut(&mut self, node_id: &Uuid) -> Option<&mut Node> {
        self.node_store_mut().nodes.get_mut(node_id)
    }

    /// Delete node from this store.
    fn del_node(&mut self, node_id: &Uuid) -> Option<Node> {
        let store = self.node_store_mut();
        let deleted = store.nodes.remove(node_id);
        if let Some(deleted) = deleted.as_ref() {
            store.aggregate.subtract(deleted.get_meta());
            store.roots.remove(node_id);
            store.leafs.remove(node_id);
            store.depth = None;
        }
        deleted
    }

    /// Update max node depth.
    fn calculate_node_depth(&self) -> usize
    where
        Self: Sized,
    {
        self.leaf_ids()
            .filter_map(|id| self.node_depth(&id, false, None, None).ok())
            .max()
            .unwrap_or(0)
    }

    /// Get all node IDs in this store.
    fn all_node_ids(&self) -> impl Iterator<Item = Uuid> {
        self.node_store().nodes.keys().copied()
    }

    /// Get all node IDs in this store.
    fn all_nodes(&self) -> impl Iterator<Item = &Node> {
        self.node_store().nodes.values()
    }

    /// Get all root node IDs in this store.
    fn root_ids(&self) -> impl Iterator<Item = Uuid> {
        self.node_store().roots.iter().copied()
    }

    /// Get all root nodes in this store.
    fn root_nodes(&self) -> impl Iterator<Item = &Node> {
        self.node_store()
            .roots
            .iter()
            .filter_map(|id| self.get_node(id))
    }

    /// Get all leaf node IDs in this store.
    fn leaf_ids(&self) -> impl Iterator<Item = Uuid> {
        self.node_store().leafs.iter().copied()
    }

    /// Get all leaf nodes in this store.
    fn leaf_nodes(&self) -> impl Iterator<Item = &Node> {
        self.node_store()
            .leafs
            .iter()
            .filter_map(|id| self.get_node(id))
    }

    /// Set a new parent value for the given child ID. Returns an error if the child
    /// node does not exist in the store. Setting the parent ID to None removes the
    /// parent-child relationship.
    fn set_parent(&mut self, child_id: &Uuid, parent_id: Option<&Uuid>) -> Result<()> {
        // Remove child from current parent.
        if let Some(child) = self.get_node(child_id) {
            if let Some(current_parent_id) = child.parent {
                let add_to_leafs = {
                    if let Some(current_parent) = self.get_node_mut(&current_parent_id) {
                        current_parent.children.retain(|x| x != child_id);
                        current_parent.is_leaf()
                    } else {
                        false
                    }
                };
                if add_to_leafs {
                    self.node_store_mut().leafs.insert(current_parent_id);
                }
            }
        } else {
            return Err(anyhow!(GraphError::NodeNotInStore(child_id.to_owned())));
        }

        // Set new parent value.
        if let Some(parent_id) = parent_id {
            if self.has_node(parent_id) {
                if let Some(child) = self.get_node_mut(child_id) {
                    child.parent = Some(parent_id.to_owned());
                    self.node_store_mut().roots.remove(child_id);
                }
                if let Some(parent) = self.get_node_mut(parent_id) {
                    parent.children.push(child_id.to_owned());
                    self.node_store_mut().leafs.remove(parent_id);
                }
            } else {
                return Err(anyhow!(GraphError::NodeNotInStore(parent_id.to_owned())));
            }
        } else if let Some(child) = self.get_node_mut(child_id) {
            child.parent = None;
            child.is_bus = false;
            self.node_store_mut().roots.insert(*child_id);
        }
        self.node_store_mut().depth = None;
        Ok(())
    }

    /// Set the is_bus value of a given node ID.
    fn set_bus(&mut self, node_id: &Uuid, is_bus: bool) -> Result<()> {
        if let Some(node) = self.get_node_mut(node_id) {
            if is_bus && node.parent.is_none() {
                return Err(anyhow!(GraphError::BusWithoutParent(node_id.to_owned())));
            }
            node.is_bus = is_bus;
        } else {
            return Err(anyhow!(GraphError::NodeNotInStore(node_id.to_owned())));
        }
        Ok(())
    }

    /// Get the bus node IDs that fall directly under the given parent ID.
    fn bus_ids(&self, parent_id: &Uuid) -> Result<impl Iterator<Item = Uuid>> {
        Ok(self
            .bus_nodes(parent_id)?
            .into_iter()
            .map(|node| node.id())
            .copied())
    }

    /// Get the bus nodes that fall directly under the given parent ID.
    fn bus_nodes(&self, parent_id: &Uuid) -> Result<impl Iterator<Item = &Node>> {
        if let Some(parent) = self.get_node(parent_id) {
            Ok(parent.children.iter().filter_map(|id| {
                self.get_node(id)
                    .and_then(|node| if node.is_bus { Some(node) } else { None })
            }))
        } else {
            Err(anyhow!(GraphError::NodeNotInStore(parent_id.to_owned())))
        }
    }

    /// Get ascendant node IDs of a given node. Setting `safe` checks for cyclic
    /// hierarchies along the way. `only_root` only includes the final root node.
    /// `root_ids` are nodes to consider as (additional) root nodes in this query.
    /// `height` determines the maximum height at which the ancestor is considered a
    /// root node.
    fn ascendant_ids<'a, 'n>(
        &'n self,
        node_id: &'a Uuid,
        safe: bool,
        only_root: bool,
        root_ids: Option<&'a BTreeSet<Uuid>>,
        height: Option<usize>,
    ) -> impl Iterator<Item = Result<&'n Uuid, Error>>
    where
        Self: Sized,
    {
        AscendantIterator::new(self, node_id, safe, only_root, root_ids, height).into_iter()
    }

    /// Get ascendant nodes of a given node. Setting `safe` checks for cyclic
    /// hierarchies along the way. `only_root` only includes the final root node.
    /// `root_ids` are nodes to consider as (additional) root nodes in this query.
    /// `height` determines the maximum height at which the ancestor is considered a
    /// root node.
    fn ascendant_nodes<'a, 'n>(
        &'n self,
        node_id: &'a Uuid,
        safe: bool,
        only_root: bool,
        root_ids: Option<&'a BTreeSet<Uuid>>,
        height: Option<usize>,
    ) -> impl Iterator<Item = Result<&'n Node, Error>>
    where
        Self: Sized,
    {
        self.ascendant_ids(node_id, safe, only_root, root_ids, height)
            .filter_map(|r| match r {
                Ok(id) => self.get_node(id).map(|n| Ok(n)),
                Err(e) => Some(Err(e)),
            })
    }

    /// Node depth (i.e. the number of levels that exist above). Setting `safe` checks
    /// for cyclic hierarchies along the way. `root_ids` are nodes to consider as (additional) root nodes in this query.
    /// `height` determines the maximum height at which the ancestor is considered a
    /// root node.
    fn node_depth(
        &self,
        node_id: &Uuid,
        safe: bool,
        root_ids: Option<&BTreeSet<Uuid>>,
        height: Option<usize>,
    ) -> Result<usize>
    where
        Self: Sized,
    {
        self.ascendant_ids(node_id, safe, false, root_ids, height)
            .try_fold(0usize, |acc, id| id.map(|_| acc + 1))
    }

    /// Get the descendant node IDs of a given node. Setting `safe` checks for cyclic
    /// hierarchies. Setting `only_leaf` only includes only absolute or specified leaf
    /// nodes in the result. `leaf_ids` restricts the search at the given node IDs,
    /// disallowing it from going any further. `depth` specifies the maximum depth at
    /// which nodes are also considered leaf nodes for this search.
    fn descendant_ids<'a, 'n>(
        &'n self,
        node_id: &'a Uuid,
        safe: bool,
        only_leaf: bool,
        leaf_ids: Option<&'a BTreeSet<Uuid>>,
        depth: Option<usize>,
    ) -> impl Iterator<Item = Result<&'n Uuid, Error>>
    where
        Self: Sized,
    {
        DescendantIterator::new(self, node_id, safe, only_leaf, leaf_ids, depth).into_iter()
    }

    /// Get the descendant nodes of a given node. Setting `safe` checks for cyclic
    /// hierarchies. Setting `only_leaf` only includes only absolute or specified leaf
    /// nodes in the result. `leaf_ids` restricts the search at the given node IDs,
    /// disallowing it from going any further. `depth` specifies the maximum depth with
    /// respect to the given node ID to search at.
    fn descendant_nodes<'a, 'n>(
        &'n self,
        node_id: &'a Uuid,
        safe: bool,
        only_leaf: bool,
        leaf_ids: Option<&'a BTreeSet<Uuid>>,
        depth: Option<usize>,
    ) -> impl Iterator<Item = Result<&'n Node, Error>>
    where
        Self: Sized,
    {
        self.descendant_ids(node_id, safe, only_leaf, leaf_ids, depth)
            .filter_map(|r| match r {
                Ok(id) => self.get_node(id).map(|n| Ok(n)),
                Err(e) => Some(Err(e)),
            })
    }

    /// Node height (i.e. the number of levels that exist below). Setting `safe` checks
    /// for cyclic hierarchies. `leaf_ids` restricts the search at the given node IDs,
    /// disallowing it from going any further. `depth` specifies the maximum depth with
    /// respect to the given node ID to search at.
    fn node_height(
        &self,
        node_id: &Uuid,
        safe: bool,
        leaf_ids: Option<&BTreeSet<Uuid>>,
        depth: Option<usize>,
    ) -> Result<usize>
    where
        Self: Sized,
    {
        let mut iterator = DescendantIterator::new(self, node_id, safe, true, leaf_ids, depth);
        while let Some(result) = iterator.next() {
            if let Err(e) = result {
                return Err(e);
            }
        }
        return Ok(iterator.levels);
    }

    /// Node width in terms of (optionally specified) leaf nodes. Setting `safe` checks
    /// for cyclic hierarchies. `leaf_ids` restricts the search at the given node IDs,
    /// disallowing it from going any further. `depth` specifies the maximum depth with
    /// respect to the given node ID to search at.
    fn node_width(
        &self,
        node_id: &Uuid,
        safe: bool,
        leaf_ids: Option<&BTreeSet<Uuid>>,
        depth: Option<usize>,
    ) -> Result<usize>
    where
        Self: Sized,
    {
        self.descendant_ids(node_id, safe, true, leaf_ids, depth)
            .try_fold(0usize, |acc, id| id.map(|_| acc + 1))
    }

    /// Get the aggregate.
    fn node_aggregate(&self) -> &Aggregate {
        &self.node_store().aggregate
    }
}

pub struct AscendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    // Iterator input data.
    /// Node store to retrieve data from.
    nodes: &'n N,
    /// Whether to check for cyclical references.
    safe: bool,
    /// Whether to only return a root node.
    only_root: bool,
    /// What nodes are considered root nodes in the search.
    root_ids: Option<&'a BTreeSet<Uuid>>,
    /// The maximum number of levels above this level to include.
    height: Option<usize>,

    // Iterator state.
    /// Last retrieved node ID.
    node_id: Uuid,
    /// Seen set of IDs.
    seen: BTreeSet<Uuid>,
    /// Whether to stop after this iteration.
    stop: bool,
    /// Processed levels.
    levels: usize,
    /// Error state, next iteration returns None.
    error: bool,
}
impl<'a, 'n, N> AscendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    pub fn new(
        nodes: &'n N,
        node_id: &'a Uuid,
        safe: bool,
        only_root: bool,
        root_ids: Option<&'a BTreeSet<Uuid>>,
        height: Option<usize>,
    ) -> Self {
        Self {
            nodes,
            safe,
            only_root,
            root_ids,
            height,
            node_id: *node_id,
            seen: BTreeSet::default(),
            stop: false,
            levels: 0,
            error: false,
        }
    }
}
impl<'a, 'n, N> Iterator for AscendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    type Item = Result<&'n Uuid, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.stop || Some(self.levels) == self.height || self.error {
            return None;
        }

        // Get the parent of the current node ID under investigation.
        let parent_id = self.nodes.get_node(&self.node_id)?.parent.as_ref()?;

        // Safety cycle check if enabled.
        if self.safe && !self.seen.insert(*parent_id) {
            self.error = true;
            return Some(Err(anyhow!(GraphError::CyclicAscendants(self.node_id))));
        }

        // The parent node under investigation.
        let parent: &'n Node = self.nodes.get_node(parent_id)?;

        // Bookkeeping for each iteration.
        self.levels += 1;

        // If we have an absolute root or node we consider to be root, stop.
        if parent.is_root()
            || self.height.map(|h| h == 1).unwrap_or(false)
            || self
                .root_ids
                .map(|x| x.contains(parent_id))
                .unwrap_or(false)
        {
            self.stop = true;
            Some(Ok(parent_id))
        } else {
            // Set the node ID under investigation to the current parent.
            self.node_id = *parent_id;

            // We have a non-root node, check if we should only return roots and recurse
            // immediately.
            if self.only_root {
                self.next()
            } else {
                Some(Ok(parent_id))
            }
        }
    }
}

pub struct DescendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    // Iterator input data.
    /// Node store to retrieve data from.
    nodes: &'n N,
    /// Whether to check for cyclical references.
    safe: bool,
    /// Whether to only return a leaf nodes.
    only_leaf: bool,
    /// What nodes are considered leaf nodes in the search.
    leaf_ids: Option<&'a BTreeSet<Uuid>>,
    /// The maximum number of levels beneath the starting node to include.
    depth: Option<usize>,

    // Iterator state.
    /// Current descendant candidates.
    descendants: Vec<Vec<Uuid>>,
    /// Seen set of IDs.
    seen: BTreeSet<Uuid>,
    /// Maximum number of levels we've seen.
    levels: usize,
    /// Error state, next iteration returns None.
    error: bool,
}
impl<'a, 'n, N> DescendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    pub fn new(
        nodes: &'n N,
        node_id: &'a Uuid,
        safe: bool,
        only_leaf: bool,
        leaf_ids: Option<&'a BTreeSet<Uuid>>,
        depth: Option<usize>,
    ) -> Self {
        let mut descendants = nodes
            .get_node(node_id)
            .map(|n| n.children.clone())
            .unwrap_or_default();
        descendants.reverse();
        let levels = if descendants.is_empty() { 0 } else { 1 };
        Self {
            nodes,
            safe,
            only_leaf,
            leaf_ids,
            depth,
            descendants: vec![descendants],
            seen: BTreeSet::from([*node_id]).into(),
            levels,
            error: false,
        }
    }
}
impl<'a, 'n, N> Iterator for DescendantIterator<'a, 'n, N>
where
    N: NodeStore,
{
    type Item = Result<&'n Uuid, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        // If we're in the error state, return None.
        if self.error {
            return None;
        }

        // Get the next node ID
        while self
            .descendants
            .last()
            .map(|ds| ds.is_empty())
            .unwrap_or_default()
        {
            self.descendants.pop();
        }
        let node_id = self.descendants.last_mut().and_then(|ds| ds.pop())?;

        // Do a safety cycle check if enabled. Get the node right after.
        if self.safe && !self.seen.insert(node_id) {
            self.error = true;
            return Some(Err(anyhow!(GraphError::CyclicDescendants(node_id))));
        }
        let node = self.nodes.get_node(&node_id)?;

        // If it is an absolute leaf or a considered leaf node or at the depth, return it.
        if node.is_leaf()
            || self
                .leaf_ids
                .map(|ids| ids.contains(&node_id))
                .unwrap_or(false)
            || self.depth.map(|d| d == self.levels).unwrap_or(false)
        {
            Some(Ok(node.id()))
        } else {
            // We have a non-leaf node, so go to try and return the next one if we want only leafs.
            // Append this node's children to the list in reversed order, so they are attempted first.
            let mut children = node.children.clone();
            children.reverse();
            self.descendants.push(children);
            self.levels = self.levels.max(self.descendants.len());
            if self.only_leaf {
                self.next()
            } else {
                Some(Ok(node.id()))
            }
        }
    }
}