// 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::collections::HashMap;
use std::num::NonZeroU16;

use anyhow::{Context, Result, bail, ensure};
use dataset_writer::arrow::array::AsArray;
use dataset_writer::arrow::datatypes::{Int32Type, UInt64Type};
use dsi_bitstream::prelude::LittleEndian;
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use rayon::prelude::*;
use swh_graph::graph::{NodeId, SwhGraph, SwhGraphWithProperties};
use swh_graph::properties;
use swh_graph_stdlib::connectivity::SubgraphWccs;
use value_traits::slices::SliceByValue;
use webgraph::prelude::*;

use crate::VulnerabilityId;

/// Converts a Parquet record batch from vuln2commit format to node-vulnerability pairs.
///
/// Reads a record batch containing `revision_id` and `vulnerability_filename` columns,
/// resolves vulnerability filenames to IDs, optionally contracts node IDs through a
/// SubgraphWccs mapping, and returns a vector of (revision_id, vulnerability_id) pairs.
fn process_batch<G: SwhGraphWithProperties<Maps: properties::Maps>>(
    num_nodes: usize,
    vulnerability_ids: &HashMap<String, VulnerabilityId>,
    graph: &G,
    subgraphwccs: Option<&SubgraphWccs<impl AsRef<[usize]>, impl SliceByValue<Value = usize>>>,
    batch: dataset_writer::arrow::array::RecordBatch,
) -> Result<Vec<Result<(NodeId, usize)>>> {
    let revision_ids = batch
        .column_by_name("revision_id")
        .context("Missing column revision_id")?
        .as_primitive_opt::<UInt64Type>()
        .context("revision_id column does not have type 'u64'")?;
    let vulnerability_filenames = batch
        .column_by_name("vulnerability_filename")
        .context("Missing column vulnerability_filename")?
        .as_dictionary_opt::<Int32Type>()
        .context("vulnerability_filenames column does not have type 'dict<i32, _>'")?;
    let vulnerability_filename_values = vulnerability_filenames
        .values()
        .as_string_opt::<i32>()
        .context("vulnerability_filenames is a dict but its values are not i32-lengthed strings")?;

    // cache the prev_vulnerability_id as it is repeated a lot due to the order in
    // which entries are written when building the input file
    let mut prev_vulnerability_filename_key = None;
    let mut prev_vulnerability_id = None;

    ensure!(
        revision_ids.len() == vulnerability_filenames.keys().len(),
        "revision_ids and vulnerability_filenames columns don't have the same number of values"
    );
    Ok(revision_ids
        .into_iter()
        .zip(vulnerability_filenames.keys())
        .map(move |(rev_id, vulnerability_filename_key)| {
        let Some(rev_id) = rev_id else {
            bail!("null value in revision_id column")
        };
        let mut rev_id = usize::try_from(rev_id).context("revision_id overflowed usize")?;

        if let Some(subgraphwccs) = subgraphwccs {
            rev_id = subgraphwccs.contraction().node_id_from_underlying(rev_id).with_context(|| format!("Revision {} ({}) has a vulnerability, but is not in the contraction of the SubgraphWccs of vulnerable commits", rev_id, graph.properties().swhid(rev_id)))?;
        }

        ensure!(
            rev_id < num_nodes,
            "revision_id {rev_id} is greater than the number of nodes in the graph ({num_nodes})"
        );
        let Some(vulnerability_filename_key) = vulnerability_filename_key else {
            bail!("null value in vulnerability_filename column");
        };
        let vulnerability_id = if Some(vulnerability_filename_key)
            == prev_vulnerability_filename_key
        {
            prev_vulnerability_id
                .expect("prev_vulnerability_id is None but prev_vulnerability_filename_key is not")
        } else {
            prev_vulnerability_filename_key = Some(vulnerability_filename_key);
            let vulnerability_filename_key = usize::try_from(vulnerability_filename_key)
                .context("vulnerability_filename_key overflows usize")?;
            let vulnerability_filename =
                vulnerability_filename_values.value(vulnerability_filename_key);
            let &vulnerability_id = vulnerability_ids.get(vulnerability_filename).context("vulnerability_filename value references vulnerability filename not in the database")?;
            let vulnerability_id =
                usize::try_from(vulnerability_id.0).expect("vulnerability row id overflowed usize");
            prev_vulnerability_id = Some(vulnerability_id);
            vulnerability_id
        };

        Ok((rev_id, vulnerability_id))
    }).collect::<Vec<_>>())
}

/// Builds a BVGraph-based commit-to-vulnerability multimap from Parquet data.
///
/// Reads vulnerability-to-commit mappings from Parquet files, inverts them to create
/// commit-to-vulnerability mappings, sorts the pairs, and writes them as a BVGraph
/// for efficient lookups.
pub fn build_commit_to_vuln_bvgraph<
    G: SwhGraph + SwhGraphWithProperties<Maps: properties::Maps> + Sync,
>(
    num_nodes: usize,
    vulnerability_ids: &HashMap<String, VulnerabilityId>,
    graph: &G,
    subgraphwccs: Option<
        &SubgraphWccs<impl AsRef<[usize]> + Sync, impl SliceByValue<Value = usize> + Sync>,
    >,
    vuln2commit_dir: &std::path::Path,
    num_partitions: NonZeroU16,
    multimap_out: &std::path::Path,
) -> Result<()> {
    log::info!("Reading Parquet file headers and counting rows");
    let reader_builders = crate::utils::open_parquet_table(vuln2commit_dir)
        .context("Could not open vuln2commit Parquet table")?;

    let num_rows: i64 = reader_builders
        .iter()
        .map(|reader_builder| reader_builder.metadata().file_metadata().num_rows())
        .sum();
    ensure!(num_rows >= 0, "Number of rows overflowed i64");
    let num_rows = usize::try_from(num_rows).context("Number of rows overflowed usize")?;

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "row",
        local_speed = true,
        expected_updates = Some(num_rows),
    );
    pl.start("[1/4] Reading and sorting vulnerable commits list");

    let unsorted_pairs = reader_builders
        .into_par_iter()
        .map(|reader_builder| {
            let mut pl = pl.clone();
            let reader = reader_builder
                .build()
                .context("Could not build Parquet reader")?;
            Ok(reader.into_iter().flat_map(move |batch| {
                match (|| -> Result<_> {
                    let batch = batch.context("Could not read record batch from file")?;
                    pl.update_with_count(batch.num_rows());
                    Ok(
                        process_batch(num_nodes, vulnerability_ids, graph, subgraphwccs, batch)?
                            .into_iter(),
                    )
                })() {
                    Ok(iter) => itertools::Either::Left(iter),
                    Err(e) => itertools::Either::Right(std::iter::once(Err(e))),
                }
            }))
        })
        .collect::<Result<Vec<_>>>()?
        .into_par_iter()
        .flatten_iter();

    let pair_sorter = webgraph::utils::par_sort_pairs::ParSortPairs::new(num_nodes)?
        .num_partitions(num_partitions.into());

    let split_iters = pair_sorter.try_sort(unsorted_pairs)?;

    let pairs: Vec<_> = split_iters.into();

    BvCompConfig::new(multimap_out).par_comp_lenders::<LittleEndian, _>(pairs, num_nodes)?;

    Ok(())
}
