use std::collections::{BTreeMap, BTreeSet};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SemanticNode {
    pub lineage_id: String,
    pub kind: String,
    pub name: String,
    pub qualified_name: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SemanticEdge {
    pub edge_id: String,
    pub kind: String,
    pub source_lineage: String,
    pub target_lineage: String,
}

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct SemanticGraph {
    nodes: BTreeMap<String, SemanticNode>,
    outgoing: BTreeMap<String, Vec<SemanticEdge>>,
    incoming: BTreeMap<String, Vec<SemanticEdge>>,
}

impl SemanticGraph {
    pub fn new(nodes: Vec<SemanticNode>, edges: Vec<SemanticEdge>) -> Self {
        let nodes_by_lineage = nodes
            .into_iter()
            .map(|node| (node.lineage_id.clone(), node))
            .collect();
        let mut graph = Self {
            nodes: nodes_by_lineage,
            outgoing: BTreeMap::new(),
            incoming: BTreeMap::new(),
        };
        for edge in edges {
            graph.add_edge(edge);
        }
        graph
    }

    pub fn add_edge(&mut self, edge: SemanticEdge) {
        self.outgoing
            .entry(edge.source_lineage.clone())
            .or_default()
            .push(edge.clone());
        self.incoming
            .entry(edge.target_lineage.clone())
            .or_default()
            .push(edge);
    }

    pub fn find_callers(&self, routine_name: &str) -> Vec<&SemanticNode> {
        let target_lineages = self.routine_lineages(routine_name);
        let caller_lineages = target_lineages
            .iter()
            .flat_map(|lineage| self.incoming.get(lineage).into_iter().flatten())
            .filter(|edge| edge.kind == "CALLS")
            .map(|edge| edge.source_lineage.as_str())
            .collect::<BTreeSet<_>>();

        caller_lineages
            .iter()
            .filter_map(|lineage| self.nodes.get(*lineage))
            .collect()
    }

    pub fn find_callees(&self, routine_name: &str) -> Vec<&SemanticNode> {
        self.routine_lineages(routine_name)
            .iter()
            .flat_map(|lineage| self.outgoing.get(lineage).into_iter().flatten())
            .filter(|edge| edge.kind == "CALLS")
            .filter_map(|edge| self.nodes.get(&edge.target_lineage))
            .collect()
    }

    pub fn find_writes(&self, routine_name: &str) -> Vec<&SemanticNode> {
        self.targets_by_edge_kind(routine_name, "WRITES")
    }

    pub fn find_reads(&self, routine_name: &str) -> Vec<&SemanticNode> {
        let query_lineages = self
            .routine_lineages(routine_name)
            .iter()
            .flat_map(|lineage| self.outgoing.get(lineage).into_iter().flatten())
            .filter(|edge| edge.kind == "OWNS_QUERY")
            .map(|edge| edge.target_lineage.as_str())
            .collect::<BTreeSet<_>>();

        query_lineages
            .iter()
            .flat_map(|lineage| self.outgoing.get(*lineage).into_iter().flatten())
            .filter(|edge| edge.kind == "READS_TABLE")
            .filter_map(|edge| self.nodes.get(&edge.target_lineage))
            .collect()
    }

    fn targets_by_edge_kind(&self, routine_name: &str, kind: &str) -> Vec<&SemanticNode> {
        self.routine_lineages(routine_name)
            .iter()
            .flat_map(|lineage| self.outgoing.get(lineage).into_iter().flatten())
            .filter(|edge| edge.kind == kind)
            .filter_map(|edge| self.nodes.get(&edge.target_lineage))
            .collect()
    }

    fn routine_lineages(&self, routine_name: &str) -> BTreeSet<String> {
        let wanted = routine_name.to_lowercase();
        self.nodes
            .values()
            .filter(|node| {
                matches!(node.kind.as_str(), "PROCEDURE" | "FUNCTION")
                    && node.name.to_lowercase() == wanted
            })
            .map(|node| node.lineage_id.clone())
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::{SemanticEdge, SemanticGraph, SemanticNode};

    #[test]
    fn resolves_callers_callees_reads_and_writes() {
        let graph = SemanticGraph::new(
            vec![
                node("routine.post", "PROCEDURE", "Проведение"),
                node("routine.check", "FUNCTION", "ПроверитьОстатки"),
                node("query.check.1", "QUERY", "ПроверитьОстатки.query1"),
                node("table.stock", "REGISTER", "ОстаткиТоваров"),
            ],
            vec![
                edge("e1", "CALLS", "routine.post", "routine.check"),
                edge("e2", "WRITES", "routine.post", "table.stock"),
                edge("e3", "OWNS_QUERY", "routine.check", "query.check.1"),
                edge("e4", "READS_TABLE", "query.check.1", "table.stock"),
            ],
        );

        assert_eq!(
            names(graph.find_callees("Проведение")),
            vec!["ПроверитьОстатки"]
        );
        assert_eq!(
            names(graph.find_callers("ПроверитьОстатки")),
            vec!["Проведение"]
        );
        assert_eq!(
            names(graph.find_writes("Проведение")),
            vec!["ОстаткиТоваров"]
        );
        assert_eq!(
            names(graph.find_reads("ПроверитьОстатки")),
            vec!["ОстаткиТоваров"]
        );
    }

    fn node(lineage_id: &str, kind: &str, name: &str) -> SemanticNode {
        SemanticNode {
            lineage_id: lineage_id.to_string(),
            kind: kind.to_string(),
            name: name.to_string(),
            qualified_name: name.to_string(),
        }
    }

    fn edge(edge_id: &str, kind: &str, source: &str, target: &str) -> SemanticEdge {
        SemanticEdge {
            edge_id: edge_id.to_string(),
            kind: kind.to_string(),
            source_lineage: source.to_string(),
            target_lineage: target.to_string(),
        }
    }

    fn names(nodes: Vec<&SemanticNode>) -> Vec<String> {
        nodes.into_iter().map(|node| node.name.clone()).collect()
    }
}