// Copyright (C) 2025-2026  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::fs::File;
use std::io::BufWriter;
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::ser::Serialize;
use mimalloc::MiMalloc;
use rapidhash::RapidBuildHasher;
use rayon::prelude::*;
use rusqlite::Connection;
use swh_graph::graph::SwhBidirectionalGraph;
use swh_graph::mph::DynMphf;
use swh_graph::views::Subgraph;
use swh_graph_stdlib::connectivity::SubgraphWccs;

use swh_osv::{db, git};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

/*
#[derive(Clone, ValueEnum, Debug, Default)]
pub enum RevisionFilter {
    #[default]
    /// All revisions connected to a revision referenced in a vulnerability report. This is the default
    All,
    /// All revisions that share a root with a revision referenced in a vulnerability report.
    SharedRoot,
}
*/

#[derive(Parser, Debug)]
#[command()]
/// Computes the connected components of each commit referenced in the vulnerability database
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 packages.rs's output
    package2vcs: Option<PathBuf>,
    #[arg(long)]
    /// Number of worker threads. Defaults to the number of CPU threads.
    parallelism: Option<NonZeroUsize>,
    #[arg(long)]
    /// If given, then any commit that shares its root directory with a vulnerable package version
    /// is included as well.
    include_vcs_matches: bool,
    #[arg(long)]
    /// Path to the epserde-serialized `SubgraphWccs` to write
    subgraphwccs_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();

    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!(
        "Listing commits in {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();
    drop(pl);

    let num_commits = commits_in_vulnerabilities.len();
    log::info!("Found {} commits", num_commits);

    if let Some(package2vcs) = &args.package2vcs {
        swh_osv::packages::graph::collect_commits_from_package2vcs(
            &graph,
            package2vcs,
            args.include_vcs_matches,
            &commits_in_vulnerabilities,
        )
        .with_context(|| format!("Could not list nodes in {}", package2vcs.display()))?;

        log::info!(
            "Found {} new nodes from packages",
            commits_in_vulnerabilities.len() - num_commits
        );
    }

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

    let wccs = SubgraphWccs::build_from_closure(
        Arc::clone(&rev_graph),
        commits_in_vulnerabilities.par_iter().map(|item| *item),
        false,
    )
    .context("Could not compute connected components")?;

    tracing::info!("Serializing SubgraphWccs");
    let mut writer = BufWriter::new(
        File::create(&args.subgraphwccs_out)
            .with_context(|| format!("Could not create {}", args.subgraphwccs_out.display()))?,
    );
    unsafe { wccs.serialize(&mut writer) }.with_context(|| {
        format!(
            "Could not write SubgraphWccs to {}",
            args.subgraphwccs_out.display()
        )
    })?;

    Ok(())
}
