// Copyright (C) 2025  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::path::Path;

use anyhow::Result;
use dsi_bitstream::prelude::LittleEndian;
use itertools::Itertools;
use lender::Lender;
use rayon::prelude::*;
use swh_graph::graph::NodeId;
use webgraph::graphs::arc_list_graph::ArcListGraph;
use webgraph::graphs::bvgraph::*;
use webgraph::prelude::MmapHelper;
use webgraph::traits::SequentialLabeling;
use webgraph::traits::graph::RandomAccessGraph;

pub trait NodeMultimap<V> {
    fn get(&self, node_id: NodeId) -> impl IntoIterator<Item = V>;
}

pub struct EmptyNodeMultimap;

impl<V> NodeMultimap<V> for EmptyNodeMultimap {
    fn get(&self, _node_id: NodeId) -> impl IntoIterator<Item = V> {
        std::iter::empty()
    }
}

pub struct BvNodeMultimap(
    pub  BvGraph<
        DynCodesDecoderFactory<
            dsi_bitstream::prelude::LE,
            MmapHelper<u32>,
            webgraph::graphs::bvgraph::EF,
        >,
    >,
);

impl BvNodeMultimap {
    pub fn mmap(path: &Path) -> Result<Self> {
        Ok(Self(
            BvGraph::with_basename(path)
                .flags(MemoryFlags::TRANSPARENT_HUGE_PAGES | MemoryFlags::RANDOM_ACCESS)
                .mode::<Mmap>()
                .endianness::<LittleEndian>()
                .load()?,
        ))
    }

    pub fn serialize(
        num_nodes: usize,
        expected_num_pairs: Option<usize>,
        path: &Path,
        pairs: impl ParallelIterator<Item = Result<(usize, usize)>>,
    ) -> Result<()> {
        let mut pair_sorter = webgraph::utils::par_sort_pairs::ParSortPairs::new(num_nodes)?;

        if let Some(expected_num_pairs) = expected_num_pairs {
            pair_sorter = pair_sorter.expected_num_pairs(expected_num_pairs);
        }

        let split_iters = pair_sorter.try_sort(pairs)?;

        let arc_list_graphs = Vec::from(split_iters.iters).into_iter().enumerate().map(
            |(partition_id, sorted_arcs_partition)| {
                ArcListGraph::new(num_nodes, sorted_arcs_partition.into_iter().dedup())
                    .iter_from(split_iters.boundaries[partition_id])
                    .take(
                        split_iters.boundaries[partition_id + 1]
                            .checked_sub(split_iters.boundaries[partition_id])
                            .expect("split_iters.boundaries is not sorted"),
                    )
            },
        );

        BvCompConfig::new(path).par_comp_lenders::<LittleEndian, _>(arc_list_graphs, num_nodes)?;

        Ok(())
    }
}

impl NodeMultimap<usize> for BvNodeMultimap {
    fn get(&self, node_id: NodeId) -> impl IntoIterator<Item = usize> {
        self.0.successors(node_id)
    }
}

/// Converts a HashMap-based multimap to a BVGraph-based multimap file.
///
/// Takes a HashMap where each key maps to a collection of values, and writes it
/// as a BVGraph file at the specified path. The BVGraph format provides efficient
/// storage and access for sparse multimaps.
pub fn serialize_hashmap_to_bvgraph(
    num_nodes: usize,
    data: &rapidhash::RapidHashMap<NodeId, impl AsRef<[NodeId]> + Default + Sync>,
    path: &Path,
) -> Result<()> {
    let max_key = data.keys().max().copied().unwrap_or(0) + 1;
    let expected_num_pairs = data.values().map(|values| values.as_ref().len()).sum();
    BvNodeMultimap::serialize(
        num_nodes.max(max_key),
        Some(expected_num_pairs),
        path,
        data.into_par_iter().flat_map(|(&key, values)| {
            values
                .as_ref()
                .into_par_iter()
                .map(move |&value| Ok((key, value)))
        }),
    )
}
