// Copyright (C) 2026  The Software Heritage developers
// See the AUTHORS file at the top-level directory of this distribution
// License: GNU General Public License version 3, or any later version
// See top-level LICENSE file for more information

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

use anyhow::Result;
use swh_graph::graph::NodeId;
use swh_graph::graph_builder::{BuiltGraph, GraphBuilder};
use swh_graph::swhid;

use swh_graph_stdlib::labeling::*;

#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct Label(BTreeSet<u32>);

impl Label {
    fn with_value(value: u32) -> Self {
        let mut set = BTreeSet::new();
        set.insert(value);
        Label(set)
    }

    fn merge(&mut self, other: &Label) {
        self.0.extend(&other.0);
    }
}

struct TestMapReducer {
    node_labels: HashMap<NodeId, u32>,
    traversed_nodes: Vec<(NodeId, Option<Label>)>,
}

impl TestMapReducer {
    fn new(node_labels: HashMap<NodeId, u32>) -> Self {
        TestMapReducer {
            node_labels,
            traversed_nodes: Vec::new(),
        }
    }

    fn run(
        self,
        graph: &BuiltGraph,
        nodes: &[NodeId],
        keep_labels: bool,
    ) -> Result<Vec<(NodeId, Option<Label>)>> {
        let mut mr = if keep_labels {
            MapReduceBuilder::new_sparse(graph, self)
                .cheap_clones(true)
                .keep_labels(true)
                .build()?
        } else {
            MapReduceBuilder::new_sparse(graph, self).build()?
        };
        mr.run_in_topological_order(nodes.iter().copied())?;
        Ok(mr.map_reducer.traversed_nodes)
    }
}

impl MapReducer for TestMapReducer {
    type Label = Label;
    type Error = anyhow::Error;

    fn map(&mut self, node: NodeId) -> Result<Option<Self::Label>> {
        Ok(self.node_labels.get(&node).map(|&v| Label::with_value(v)))
    }

    fn reduce<'a, I: Iterator<Item = &'a Self::Label>>(
        &mut self,
        first_label: Self::Label,
        other_labels: I,
    ) -> Result<Option<Self::Label>, Self::Error>
    where
        Self::Label: 'a,
    {
        let result = other_labels.fold(first_label, |mut left, right| {
            left.merge(right);
            left
        });
        Ok(if result.0.is_empty() {
            None
        } else {
            Some(result)
        })
    }

    fn on_node_traversed(&mut self, node: NodeId, label: Option<&Self::Label>) -> Result<()> {
        self.traversed_nodes.push((node, label.cloned()));
        Ok(())
    }
}

fn build_linear_graph() -> Result<(BuiltGraph, Vec<NodeId>)> {
    // 0 <- 1 <- 2 <- 3
    let mut builder = GraphBuilder::default();
    let n0 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000000))?
        .done();
    let n1 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000001))?
        .done();
    let n2 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000002))?
        .done();
    let n3 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000003))?
        .done();

    builder.arc(n1, n0);
    builder.arc(n2, n1);
    builder.arc(n3, n2);

    Ok((builder.done()?, vec![n0, n1, n2, n3]))
}

#[test]
fn test_linear_chain() -> Result<()> {
    let (graph, nodes) = build_linear_graph()?;
    let [n0, n1, n2, n3] = nodes[..] else {
        panic!("Expected 4 nodes");
    };

    let mut labels = HashMap::new();
    labels.insert(n0, 10);
    labels.insert(n2, 20);

    let expected = HashMap::from([
        (n0, Some(Label(BTreeSet::from([10])))),
        (n1, Some(Label(BTreeSet::from([10])))),
        (n2, Some(Label(BTreeSet::from([10, 20])))),
        (n3, Some(Label(BTreeSet::from([10, 20])))),
    ]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}

#[test]
fn test_node_with_no_label() -> Result<()> {
    let (graph, nodes) = build_linear_graph()?;
    let [n0, n1, n2, n3] = nodes[..] else {
        panic!("Expected 4 nodes");
    };

    let labels = HashMap::new();

    let expected = HashMap::from([(n0, None), (n1, None), (n2, None), (n3, None)]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}

fn build_diamond_graph() -> Result<(BuiltGraph, Vec<NodeId>)> {
    //     0
    //     ^
    //    / \
    //   1   2
    //   ^   ^
    //    \ /
    //     3
    let mut builder = GraphBuilder::default();
    let n0 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000000))?
        .done();
    let n1 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000001))?
        .done();
    let n2 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000002))?
        .done();
    let n3 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000003))?
        .done();

    builder.arc(n1, n0);
    builder.arc(n2, n0);
    builder.arc(n3, n1);
    builder.arc(n3, n2);

    Ok((builder.done()?, vec![n0, n1, n2, n3]))
}

#[test]
fn test_diamond_graph() -> Result<()> {
    let (graph, nodes) = build_diamond_graph()?;
    let [n0, n1, n2, n3] = nodes[..] else {
        panic!("Expected 4 nodes");
    };

    let mut labels = HashMap::new();
    labels.insert(n0, 5);
    labels.insert(n1, 10);
    labels.insert(n2, 20);

    let expected = HashMap::from([
        (n0, Some(Label(BTreeSet::from([5])))),
        (n1, Some(Label(BTreeSet::from([5, 10])))),
        (n2, Some(Label(BTreeSet::from([5, 20])))),
        (n3, Some(Label(BTreeSet::from([5, 10, 20])))),
    ]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}

#[test]
fn test_callback_receives_all_nodes() -> Result<()> {
    let (graph, nodes) = build_diamond_graph()?;

    let mut labels = HashMap::new();
    labels.insert(nodes[0], 1);

    let traversed_set: BTreeSet<NodeId> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .map(|(n, _)| n)
        .collect();
    assert_eq!(traversed_set, nodes.iter().copied().collect());

    let traversed_set: BTreeSet<NodeId> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .map(|(n, _)| n)
        .collect();
    assert_eq!(traversed_set, nodes.iter().copied().collect());

    Ok(())
}

fn build_tree_graph() -> Result<(BuiltGraph, Vec<NodeId>)> {
    //       0
    //       ^
    //      / \
    //     1   2
    //     ^
    //    / \
    //   3   4
    let mut builder = GraphBuilder::default();
    let n0 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000000))?
        .done();
    let n1 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000001))?
        .done();
    let n2 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000002))?
        .done();
    let n3 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000003))?
        .done();
    let n4 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000004))?
        .done();

    builder.arc(n1, n0);
    builder.arc(n2, n0);
    builder.arc(n3, n1);
    builder.arc(n4, n1);

    Ok((builder.done()?, vec![n0, n1, n3, n4, n2]))
}

#[test]
fn test_tree_graph() -> Result<()> {
    let (graph, nodes) = build_tree_graph()?;
    let [n0, n1, n3, n4, n2] = nodes[..] else {
        panic!("Expected 5 nodes");
    };

    let mut labels = HashMap::new();
    labels.insert(n0, 1);
    labels.insert(n3, 30);
    labels.insert(n4, 40);

    let expected = HashMap::from([
        (n0, Some(Label(BTreeSet::from([1])))),
        (n1, Some(Label(BTreeSet::from([1])))),
        (n2, Some(Label(BTreeSet::from([1])))),
        (n3, Some(Label(BTreeSet::from([1, 30])))),
        (n4, Some(Label(BTreeSet::from([1, 40])))),
    ]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}

fn build_mixed_label_graph() -> Result<(BuiltGraph, Vec<NodeId>)> {
    //   0   2   4
    //   ^   ^   ^
    //   |   |   |
    //   1   6   |
    //   ^   ^   |
    //    \ /    |
    //     3    /
    //     ^   /
    //      \ /
    //       5
    // where only n2 and n4 are labeled (so n1 also doesn't get a label)
    let mut builder = GraphBuilder::default();
    let n0 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000000))?
        .done();
    let n1 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000001))?
        .done();
    let n2 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000002))?
        .done();
    let n3 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000003))?
        .done();
    let n4 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000004))?
        .done();
    let n5 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000005))?
        .done();

    builder.arc(n1, n0);
    builder.arc(n3, n1);
    builder.arc(n3, n2);
    builder.arc(n5, n3);
    builder.arc(n5, n4);

    Ok((builder.done()?, vec![n0, n1, n2, n3, n4, n5]))
}

#[test]
fn test_mixed_label_presence() -> Result<()> {
    let (graph, nodes) = build_mixed_label_graph()?;
    let [n0, n1, n2, n3, n4, n5] = nodes[..] else {
        panic!("Expected 6 nodes");
    };

    let mut labels = HashMap::new();
    labels.insert(n2, 20);
    labels.insert(n4, 40);

    let expected = HashMap::from([
        (n0, None),
        (n1, None),
        (n2, Some(Label(BTreeSet::from([20])))),
        (n3, Some(Label(BTreeSet::from([20])))),
        (n4, Some(Label(BTreeSet::from([40])))),
        (n5, Some(Label(BTreeSet::from([20, 40])))),
    ]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}

#[test]
fn test_node_without_successors() -> Result<()> {
    let mut builder = GraphBuilder::default();
    let n0 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000000))?
        .done();
    let n1 = builder
        .node(swhid!(swh:1:rev:0000000000000000000000000000000000000001))?
        .done();
    builder.arc(n1, n0);
    let graph = builder.done()?;
    let nodes = vec![n0, n1];

    let mut labels = HashMap::new();
    labels.insert(n0, 42);

    let expected = HashMap::from([
        (n0, Some(Label(BTreeSet::from([42])))),
        (n1, Some(Label(BTreeSet::from([42])))),
    ]);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels.clone())
        .run(&graph, &nodes, false)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    let traversed: HashMap<_, _> = TestMapReducer::new(labels)
        .run(&graph, &nodes, true)?
        .into_iter()
        .collect();
    assert_eq!(traversed, expected);

    Ok(())
}
