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

use anyhow::{Context, Result};
use clap::Parser;
use dashmap::DashSet;
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger, progress_logger};
use epserde::deser::{Deserialize, Flags};
use mimalloc::MiMalloc;
use rapidhash::RapidBuildHasher;
use rayon::prelude::*;
use rusqlite::Connection;
use swh_graph::graph::{SwhBidirectionalGraph, SwhGraph, SwhGraphWithProperties};
use swh_graph::mph::DynMphf;
use swh_graph::views::Subgraph;
use swh_graph_stdlib::connectivity::SubgraphWccs;

use swh_osv::{cherrypicks, db, git};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
#[command()]
/// Reads the vulnerability database, looks for cherry-picks of any commit mentioned in it, and
/// writes them as a Parquet table
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)]
    /// Number of worker threads. Defaults to the number of CPU threads.
    parallelism: Option<NonZeroUsize>,
    #[arg(long)]
    /// If given, strips redundant cherrypick links
    ///
    /// For example, if A is a cherry-pick of B and C, and B is a cherry-pick of C,
    /// then --reduce-transitive-cherrypicks omits the information that A is a cherry-pick of C
    /// from the output
    reduce_transitive_cherrypicks: bool,
    #[arg(long)]
    /// If given, written node ids are those in the ContiguousSubgraph made of commits
    /// in the connected components, not in the whole graph
    subgraph_nodeids: bool,
    #[arg(long)]
    /// Path to the epserde-serialized `SubgraphWccs` written by `compute-connected-components`
    subgraphwccs: PathBuf,
    #[arg(long)]
    /// Path to the epserde-serialized cherrypick-to-original multimap to write
    cherrypicks2originals_out: Option<PathBuf>,
    #[arg(long)]
    /// Path to the epserde-serialized original-to-cherrypick multimap to write
    originals2cherrypicks_out: Option<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();

    tracing::debug!("Opening sqlite database...");
    let conn = Connection::open(&args.db)
        .with_context(|| format!("Could not open {}", args.db.display()))?;

    tracing::info!("Loading graph...");
    let graph = SwhBidirectionalGraph::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")?
        .load_properties(|properties| properties.load_strings())
        .context("Could not load strings")?;
    let graph = Arc::new(graph);

    tracing::info!("Counting vulnerabilities...");
    let num_vulnerabilities =
        db::count_vulnerabilities(&conn).context("Could not count vulnerabilities")?;

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(num_vulnerabilities),
    );
    pl.threshold(1000);
    pl.start(format!(
        "Processing {num_vulnerabilities} vulnerabilities..."
    ));

    // set of all nodes in any of the connected components of a commit mentioned by the
    // vulnerability database
    let commits_in_vulnerabilities = DashSet::with_hasher(RapidBuildHasher::default());

    let git_range_parser = git::GitRangeParser::new(Arc::clone(&graph));

    db::collect_commits_from_vulnerability_ranges(
        &conn,
        &pl,
        args.parallelism,
        &git_range_parser,
        &commits_in_vulnerabilities,
    )?;

    pl.done();

    let graph = Arc::new(Subgraph::with_node_constraint(
        Arc::clone(&graph),
        "rev".parse().unwrap(),
    ));

    tracing::info!("Loading SubgraphWccs...");
    let wccs = unsafe { <SubgraphWccs>::mmap(&args.subgraphwccs, Flags::RANDOM_ACCESS) }
        .with_context(|| {
            format!(
                "Could not mmap connected components from {}",
                args.subgraphwccs.display()
            )
        })?;
    let wccs = wccs.uncase();

    let num_connected_nodes = wccs.num_nodes();
    let num_connected_components = wccs.num_components();
    let (stats, mut cherrypicks_to_originals) =
        cherrypicks::list_cherrypicks(&Arc::clone(&graph), wccs)?;
    let num_nodes_in_cherrypick_graph = if args.subgraph_nodeids {
        let mut pl = progress_logger!(
            display_memory = true,
            item_name = "cherrypick",
            local_speed = true,
            expected_updates = Some(cherrypicks_to_originals.0.len()),
        );
        pl.start("Converting to compacted subgraph node ids");
        cherrypicks_to_originals.0 = cherrypicks_to_originals
            .0
            .into_par_iter()
            .map_with(pl.clone(), |pl, (cherrypick, originals)| {
                pl.light_update();
                let Some(cherrypick_contracted_id) =
                    wccs.contraction().node_id_from_underlying(cherrypick)
                else {
                    panic!(
                        "{} (cherry-pick of {}) is not in the subgraph",
                        graph.properties().swhid(cherrypick),
                        originals
                            .iter()
                            .map(|&original| format!("{}", graph.properties().swhid(original)))
                            .next()
                            .unwrap_or_else(|| "<none???>".to_string()),
                    )
                };
                (
                    cherrypick_contracted_id,
                    originals
                        .into_iter()
                        .map(|original| {
                            let Some(original_contracted_id) =
                                wccs.contraction().node_id_from_underlying(original)
                            else {
                                panic!(
                                    "{} (cherry-picked by {}) is not in the subgraph",
                                    graph.properties().swhid(original),
                                    graph.properties().swhid(cherrypick)
                                )
                            };
                            original_contracted_id
                        })
                        .collect(),
                )
            })
            .collect();
        pl.done();
        wccs.contraction().num_nodes()
    } else {
        graph.num_nodes()
    };
    let num_cherrypicks: usize = cherrypicks_to_originals.0.values().map(Vec::len).sum();
    let originals_to_cherrypicks = cherrypicks::Original2Cherrypicks(
        cherrypicks::invert_cherrypick_map(
            cherrypicks_to_originals.0.len(),
            &cherrypicks_to_originals.0,
        )
        .collect(),
    );
    let num_cherrypicked: usize = originals_to_cherrypicks.0.values().map(Vec::len).sum();

    let swh_osv::cherrypicks::Stats {
        no_matches,
        no_matches_but_mentions_cherrypick,
        at_least_one_match,
        at_least_two_matches,
        unparsable_commit: _,
        unknown_commit_id_and_no_repo_url,
        unknown_commit_id_and_has_nonutf8_origin_url,
        unknown_commit_id_and_has_known_origin_url,
        unknown_commit_id_and_has_unknown_origin_url,
        different_component,
        found,
    } = stats;
    println!(
        "{num_connected_nodes:?} commits are in the {num_connected_components:?} connected components that contain any of the `introduced`, `fixed`, `last_affected`, or `limit` commits mentioned by a vulnerability reports.\n"
    );
    println!("Of the commits mentioned by a vulnerability reports:");
    println!(
        "* {no_matches:?} are not deemed to be cherry-picks, though {no_matches_but_mentions_cherrypick:?} mention the keywords \"cherry picked from commit\" (this typically happens because a cherry-picked commit's message is quoted in an other commit message)"
    );
    println!(
        "* {at_least_one_match:?} have at least one valid \"cherry picked from commit\" stanza, {at_least_two_matches:?} have at least two. Of those stanzas:"
    );
    println!(
        "    * {unknown_commit_id_and_no_repo_url:?} reference unknown commits and don't have a repo URL"
    );
    println!(
        "    * {unknown_commit_id_and_has_nonutf8_origin_url:?} reference unknown commits and a non-UTF8 repo URL"
    );
    println!(
        "    * {unknown_commit_id_and_has_known_origin_url:?} reference unknown repo URLs (ie. these were origins unknown to SWH at the time of graph export)"
    );
    println!(
        "    * {unknown_commit_id_and_has_unknown_origin_url:?} reference unknown commits and a known repo URL (ie. these commits were unknown at the time of the export, but probably will be in the future)"
    );
    println!(
        "    * {different_component:?} claim to be cherry-picks of commits in different connected components"
    );
    println!("    * {found:?} reference known commits");
    println!("* {num_cherrypicks:?} are cherry-picks of known commits");
    println!("* {num_cherrypicked:?} are cherry-picked by any commit");

    let (reduction_stats, reduced_cherrypicks_to_originals) =
        cherrypicks::reduce_transitive_cherrypicks(&*graph, &cherrypicks_to_originals.0);

    println!(
        "Of the {} cherry-pick commits, {} are cherry-picks of multiple commits.",
        reduction_stats.total_cherrypicks, reduction_stats.total_nonsingleton_cherrypicks
    );
    println!(
        "Of the latter, {} can be reduced to cherry-picks of a single commit, which is itself a cherry-pick of one (or more). {} can only be partially reduced through that process.",
        reduction_stats.reducible_cherrypickedfrom,
        reduction_stats.partially_reducible_cherrypickedfrom
    );

    let reduced_cherrypicks_to_originals = if args.reduce_transitive_cherrypicks {
        Some(reduced_cherrypicks_to_originals)
    } else {
        None
    };

    if let Some(cherrypicks2originals_out) = args.cherrypicks2originals_out {
        let mut pl = progress_logger!(
            display_memory = true,
            item_name = "node",
            local_speed = true,
            expected_updates = Some(graph.num_nodes()),
        );
        pl.start("Building cherrypick-to-original multimap");
        let map = reduced_cherrypicks_to_originals
            .as_ref()
            .unwrap_or(&cherrypicks_to_originals.0);
        swh_osv::multimap::serialize_hashmap_to_bvgraph(
            num_nodes_in_cherrypick_graph,
            map,
            &cherrypicks2originals_out,
        )?;
    }

    if let Some(originals2cherrypicks_out) = args.originals2cherrypicks_out {
        if let Some(reduced_cherrypicks_to_originals) = reduced_cherrypicks_to_originals {
            let reduced_originals_to_cherrypicks = cherrypicks::invert_cherrypick_map(
                reduced_cherrypicks_to_originals.len(),
                reduced_cherrypicks_to_originals,
            )
            .collect();
            swh_osv::multimap::serialize_hashmap_to_bvgraph(
                num_nodes_in_cherrypick_graph,
                &reduced_originals_to_cherrypicks,
                &originals2cherrypicks_out,
            )?;
        } else {
            swh_osv::multimap::serialize_hashmap_to_bvgraph(
                num_nodes_in_cherrypick_graph,
                &originals_to_cherrypicks.0,
                &originals2cherrypicks_out,
            )?;
        }
    }

    tracing::info!("Done.");

    Ok(())
}
