// 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::num::NonZeroU16;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::Parser;
use epserde::prelude::{Deserialize, Flags};
//use mimalloc::MiMalloc;
use swh_graph::graph::{SwhGraph, SwhUnidirectionalGraph};
use swh_graph::mph::DynMphf;
use swh_graph_stdlib::connectivity::SubgraphWccs;

// mimalloc makes the whole script twice as slow as GNU malloc
// (and it requests over 100TB of virtual memory for no reason with the default config
// and 96 threads, but it doesn't seem to be an issue in practice)
//#[global_allocator]
//static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
#[command()]
/// Reads a Parquet table written by build-vuln2commit-parquet that contains pairs of
/// (revision_id, vulnerability_filename) and turns it into a BVGraph
/// that gives quick access to the set of vulnerability ids for each revision id.
///
/// Memory usage is in the order of RAYON_NUM_THREADS * sort_batch_size * num_partitions * 16 bytes,
/// ie. 800GB on a machine with 96 CPU threads with the default config.
///
/// Memory is allocated quickly at the beginning and should lower over time.
struct Args {
    #[arg(long)]
    /// Path to swh-graph's files
    graph: PathBuf,
    #[arg(long)]
    /// Path to swh-osv's sqlite database
    db: PathBuf,
    #[arg(long)]
    /// Path to the epserde-serialized `SubgraphWccs` written by `compute-connected-components`
    subgraphwccs: Option<PathBuf>,
    #[arg(long)]
    /// Path to the vulnerable commits database written by list-vulnerable-commits
    vuln2commit_dir: PathBuf,
    #[arg(long, default_value_t = NonZeroU16::new(256).unwrap())]
    /// How many partitions the range of commit ids should be split into before merging.
    ///
    /// Larger values make the final single-threaded merge step logarithmically faster,
    /// but use proportionally more RAM and file descriptors during the initial sort,
    /// and files during the whole runtime.
    num_partitions: NonZeroU16,
    #[arg(long)]
    /// Where to write the commit id -> vulnerability id map
    multimap_out: PathBuf,
}

pub fn main() -> Result<()> {
    let args = Args::parse();

    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive(tracing::Level::INFO.into()),
        )
        .init();

    let graph = SwhUnidirectionalGraph::new(&args.graph)
        .with_context(|| format!("Could not load graph from {}", args.graph.display()))?
        .init_properties()
        .load_properties(|properties| properties.load_maps::<DynMphf>())
        .context("Could not load maps")?;
    let graph = Arc::new(graph);

    let subgraphwccs = args
        .subgraphwccs
        .as_ref()
        .map(|path| -> Result<_> {
            tracing::info!("Loading SubgraphWccs...");

            unsafe { <SubgraphWccs>::mmap(path, Flags::RANDOM_ACCESS) }.with_context(|| {
                format!(
                    "Could not mmap connected components from {}",
                    path.display()
                )
            })
        })
        .transpose()?;
    let subgraphwccs = subgraphwccs.as_ref().map(|wccs| wccs.uncase());

    let num_nodes = subgraphwccs
        .as_ref()
        .map(|subgraphwccs| subgraphwccs.num_nodes())
        .unwrap_or(graph.num_nodes());

    let vulnerability_ids = swh_osv::db::read_filenames_from_sqlite(&args.db)
        .context("Could not read vulnerability filenames from sqlite")?;

    swh_osv::commit2vuln_parquet_to_bvgraph::build_commit_to_vuln_bvgraph(
        num_nodes,
        &vulnerability_ids,
        &graph,
        subgraphwccs,
        &args.vuln2commit_dir,
        args.num_partitions,
        &args.multimap_out,
    )
}
