// 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::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};
use epserde::deser::Deserialize;
use epserde::prelude::Flags;
use itertools::Itertools;
use mimalloc::MiMalloc;
use rapidhash::RapidBuildHasher;
use rusqlite::Connection;
use swh_graph::graph::{SwhBidirectionalGraph, SwhGraph};
use swh_graph::mph::DynMphf;
use swh_graph_stdlib::connectivity::SubgraphWccs;

use swh_osv::multimap::BvNodeMultimap;
use swh_osv::{Stats, db, git, reference_metrics, vuln2commit_parquet};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
/// Writes an exhaustive list of all commits between an "introduced" commit and a
/// "fixed"/"limit"/"last_affected" commit.
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)]
    /// Path to the epserde-serialized `SubgraphWccs` written by `compute-connected-components`
    subgraphwccs: Option<PathBuf>,
    #[arg(long)]
    /// Path to list of cherry-picks extracted by the 'list-cherry-picks' command
    ///
    /// If provided, the 'introduced' and 'fixed' lists are extended to include
    /// the transitive closure by cherry-picking (ie. if A is a cherry-pick of B
    /// and B fixes a vulnerability, then we assume A fixes the vulnerability too)
    cherrypicks2originals: Option<PathBuf>,
    #[arg(long)]
    /// Path to list of cherry-picks extracted by the 'list-cherry-picks' command
    ///
    /// If provided, the 'introduced' and 'fixed' lists are extended to include
    /// the transitive closure by cherry-picking (ie. if A is a cherry-pick of B
    /// and B fixes a vulnerability, then we assume A fixes the vulnerability too)
    originals2cherrypicks: Option<PathBuf>,
    #[arg(long)]
    /// Path to the Parquet database to write
    out_dir: Option<PathBuf>,
    #[arg(long)]
    /// path to a Metrics.json file as generated by the authors of the
    /// "Chasing One-day Vulnerabilities Across Open Source Forks" paper
    reference_metrics: 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()))?;
    let reference_metrics = args
        .reference_metrics
        .map(|reference_metrics_path| -> Result<_> {
            tracing::info!("Loading reference metrics");
            let reader = std::io::BufReader::new(
                std::fs::File::open(&reference_metrics_path).with_context(|| {
                    format!("Could not open {}", reference_metrics_path.display())
                })?,
            );
            let mut by_vuln_id = HashMap::new();
            let metrics: Vec<reference_metrics::ReferenceMetric> =
                serde_json::from_reader(reader).context("Could not parse reference metrics")?;
            for metric in metrics {
                by_vuln_id
                    .entry(metric.range.vulnerability_id.clone())
                    .or_insert(Vec::new())
                    .push(metric);
            }
            Ok(by_vuln_id)
        })
        .transpose()?;
    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")?;
    let graph = Arc::new(graph);

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

    let subgraphwccs = args
        .subgraphwccs
        .as_ref()
        .map(|path| -> Result<_> {
            if args.originals2cherrypicks.is_none() && args.cherrypicks2originals.is_none() {
                tracing::warn!(
                    "--subgraphwccs has no effect when --originals2cherrypicks and --cherrypicks2originals are not given"
                );
            }
            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 originals_to_cherrypicks = args.originals2cherrypicks.as_ref().map(|path| {
        if args.subgraphwccs.is_none() {
            tracing::warn!("--originals2cherrypicks given without --subgraphwccs. Assuming that commit ids are not reduced to the connected components subgraph (ie. that the 'list-cherry-picks' command was called without --subgraph-nodeids).");
        }
        tracing::info!("Loading Original2Cherrypicks...");
        BvNodeMultimap::mmap(path)
    }).transpose()?;

    let cherrypicks_to_originals = args.cherrypicks2originals.as_ref().map(|path| {
        if args.subgraphwccs.is_none() {
            tracing::warn!("--cherrypicks2originals given without --subgraphwccs. Assuming that commit ids are not reduced to the connected components subgraph (ie. that the 'list-cherry-picks' command was called without --subgraph-nodeids).");
        }
        tracing::info!("Loading Cherrypick2Originals...");
        BvNodeMultimap::mmap(path)
    }).transpose()?;

    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..."
    ));

    let stats = Stats::default();

    let all_vulnerable_commits: DashSet<usize, _> =
        DashSet::with_capacity_and_hasher(10_000_000, RapidBuildHasher::default());

    let dataset_writer = args
        .out_dir
        .map(|out_dir| {
            dataset_writer::ParallelDatasetWriter::<
                dataset_writer::ParquetTableWriter<vuln2commit_parquet::Builder>,
            >::with_schema(
                out_dir.clone(),
                (
                    Arc::new(vuln2commit_parquet::schema()),
                    vuln2commit_parquet::writer_properties(&graph).build(),
                ),
            )
            .with_context(|| format!("Could not create {}", out_dir.display()))
        })
        .transpose()?;

    let contraction = subgraphwccs
        .as_ref()
        .map(|subgraphwccs| subgraphwccs.contraction());

    let git_range_parser = git::GitRangeParser::new(Arc::clone(&graph))
        .with_nodeid_contraction(contraction)
        .with_context(|| {
            format!(
                "Could not use Contraction from {} with graph {}",
                args.subgraphwccs
                    .expect("GitRangeParser::with_nodeid_contraction(None) failed")
                    .display(),
                graph.path().display()
            )
        })?
        .with_node_multimaps::<_, _>(originals_to_cherrypicks, cherrypicks_to_originals)
        .with_context(|| {
            format!(
                "Could not use multimaps {} and {} with graph {}",
                args.originals2cherrypicks
                    .map(|path| path.display().to_string())
                    .unwrap_or("<none>".into()),
                args.cherrypicks2originals
                    .map(|path| path.display().to_string())
                    .unwrap_or("<none>".into()),
                graph.path().display()
            )
        })?;

    db::for_each_vulnerability(
        &conn,
        &pl,
        args.parallelism,
        || {
            dataset_writer
                .as_ref()
                .map(|dataset_writer| dataset_writer.get_thread_writer().unwrap())
        },
        |writer, _rowid, filename, doc| {
            if let Some(vulnerable_commits) =
                vuln2commit_parquet::list_commits_affected_by_vulnerability(
                    filename.clone(),
                    doc,
                    &git_range_parser,
                    &stats,
                    &reference_metrics,
                )?
            {
                for &commit in &vulnerable_commits {
                    all_vulnerable_commits.insert(commit);
                }
                if let Some(writer) = writer {
                    let builder = writer.builder()?;
                    builder
                        .vulnerability_filename
                        .append_values(&filename, vulnerable_commits.len());
                    for commit in vulnerable_commits.into_iter().sorted() {
                        builder
                            .revision_id
                            .append_value(u64::try_from(commit).expect("NodeId overflows u64"));
                    }
                }
            }
            Ok(())
        },
    )?;

    pl.done();

    println!("Statistics:\n{stats:#?}");
    println!("Total vulnerable_commits: {}", all_vulnerable_commits.len());

    Ok(())
}
