use std::collections::{HashMap, HashSet, VecDeque};
use std::error::Error;
use std::fmt::{Display, Formatter};

use crate::model::{Edge, Graph};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GraphFilterError {
    RootNotFound { root_iri: String },
    NoDescendants { root_iri: String },
}

impl Display for GraphFilterError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            GraphFilterError::RootNotFound { root_iri } => {
                write!(
                    f,
                    "root class IRI {root_iri} was not found in the imported graph"
                )
            }
            GraphFilterError::NoDescendants { root_iri } => {
                write!(
                    f,
                    "root class IRI {root_iri} has no subclass descendants in the imported graph"
                )
            }
        }
    }
}

impl Error for GraphFilterError {}

pub fn filter_graph_to_descendants(
    graph: &Graph,
    root_iri: &str,
) -> Result<Graph, GraphFilterError> {
    let Some(root_index) = graph
        .nodes
        .iter()
        .position(|node| node.label.as_deref() == Some(root_iri))
    else {
        return Err(GraphFilterError::RootNotFound {
            root_iri: root_iri.to_owned(),
        });
    };

    let mut adjacency = vec![Vec::new(); graph.nodes.len()];
    for edge in &graph.edges {
        adjacency[edge.source].push(edge.target);
    }

    let mut visited = HashSet::from([root_index]);
    let mut queue = VecDeque::from([root_index]);
    while let Some(node) = queue.pop_front() {
        for &child in &adjacency[node] {
            if visited.insert(child) {
                queue.push_back(child);
            }
        }
    }

    if visited.len() <= 1 {
        return Err(GraphFilterError::NoDescendants {
            root_iri: root_iri.to_owned(),
        });
    }

    let mut reindex = HashMap::new();
    let mut nodes = Vec::new();
    for (old_index, node) in graph.nodes.iter().enumerate() {
        if visited.contains(&old_index) {
            let new_index = nodes.len();
            reindex.insert(old_index, new_index);
            nodes.push(node.clone());
        }
    }

    let mut edges = Vec::new();
    for edge in &graph.edges {
        if visited.contains(&edge.source) && visited.contains(&edge.target) {
            edges.push(Edge::new(reindex[&edge.source], reindex[&edge.target]));
        }
    }

    Ok(Graph::new(nodes, edges))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::Node;

    fn node(label: &str) -> Node {
        Node {
            label: Some(label.to_owned()),
            ..Node::default()
        }
    }

    #[test]
    fn keeps_root_and_all_descendants() {
        let graph = Graph::new(
            vec![
                node("root"),
                node("child"),
                node("grandchild"),
                node("other"),
                node("other_child"),
            ],
            vec![Edge::new(0, 1), Edge::new(1, 2), Edge::new(3, 4)],
        );

        let filtered = filter_graph_to_descendants(&graph, "root").unwrap();

        assert_eq!(filtered.nodes.len(), 3);
        assert_eq!(
            filtered
                .nodes
                .iter()
                .map(|node| node.label.clone())
                .collect::<Vec<_>>(),
            vec![
                Some("root".to_owned()),
                Some("child".to_owned()),
                Some("grandchild".to_owned()),
            ]
        );
        assert_eq!(filtered.edges, vec![Edge::new(0, 1), Edge::new(1, 2)]);
    }

    #[test]
    fn returns_error_when_root_is_missing() {
        let graph = Graph::new(vec![node("other")], vec![]);

        let error = filter_graph_to_descendants(&graph, "root").unwrap_err();

        assert_eq!(
            error,
            GraphFilterError::RootNotFound {
                root_iri: "root".to_owned(),
            }
        );
    }

    #[test]
    fn returns_error_when_root_has_no_descendants() {
        let graph = Graph::new(vec![node("root"), node("other")], vec![Edge::new(1, 0)]);

        let error = filter_graph_to_descendants(&graph, "root").unwrap_err();

        assert_eq!(
            error,
            GraphFilterError::NoDescendants {
                root_iri: "root".to_owned(),
            }
        );
    }
}
