// 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(),
        }
    }
}

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 node_labels = HashMap::new();
    node_labels.insert(n0, 10);
    node_labels.insert(n2, 20);

    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(nodes.into_iter())?;

    let traversed: HashMap<NodeId, Option<Label>> =
        mr.map_reducer.traversed_nodes.iter().cloned().collect();

    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])))),
    ]);
    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 node_labels = HashMap::new();
    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(nodes.into_iter())?;

    let traversed: HashMap<NodeId, Option<Label>> =
        mr.map_reducer.traversed_nodes.iter().cloned().collect();

    let expected = HashMap::from([(n0, None), (n1, None), (n2, None), (n3, None)]);
    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 node_labels = HashMap::new();
    node_labels.insert(n0, 5);
    node_labels.insert(n1, 10);
    node_labels.insert(n2, 20);

    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(nodes.into_iter())?;

    let traversed: HashMap<NodeId, Option<Label>> =
        mr.map_reducer.traversed_nodes.iter().cloned().collect();

    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])))),
    ]);
    assert_eq!(traversed, expected);

    Ok(())
}

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

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

    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(nodes.iter().copied())?;

    let traversed_set: BTreeSet<NodeId> = mr
        .map_reducer
        .traversed_nodes
        .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 node_labels = HashMap::new();
    node_labels.insert(n0, 1);
    node_labels.insert(n3, 30);
    node_labels.insert(n4, 40);

    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(nodes.into_iter())?;

    let traversed: HashMap<NodeId, Option<Label>> =
        mr.map_reducer.traversed_nodes.iter().cloned().collect();

    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])))),
    ]);
    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 mut node_labels = HashMap::new();
    node_labels.insert(n0, 42);

    let map_reducer = TestMapReducer::new(node_labels);
    let mut mr = MapReduce::new(graph, map_reducer);
    mr.run_in_topological_order(vec![n0, n1].into_iter())?;

    let traversed: HashMap<NodeId, Option<Label>> =
        mr.map_reducer.traversed_nodes.iter().cloned().collect();

    let expected = HashMap::from([
        (n0, Some(Label(BTreeSet::from([42])))),
        (n1, Some(Label(BTreeSet::from([42])))),
    ]);
    assert_eq!(traversed, expected);
    Ok(())
}
