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

use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};
use dashmap::DashSet;
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use mimalloc::MiMalloc;
use rapidhash::RapidBuildHasher;
use rayon::prelude::*;
use rusqlite::Connection;
use swh_graph::NodeType;
use swh_graph::graph::{
    SwhBackwardGraph, SwhBidirectionalGraph, SwhForwardGraph, SwhGraphWithProperties,
};
use swh_graph::mph::DynMphf;
use swh_graph::views::Subgraph;

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(Clone, ValueEnum, Debug)]
pub enum OriginFilter {
    /// All origins connected to a revision referenced in a vulnerability report. This is the default,
    /// and returns more origins than other filters
    Connected,
    /// All origins that ever contained a revision referenced in a vulnerability report
    Contains,
}

#[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)]
    /// If --origins-out is given, limits the number of origins written to this number.
    ///
    /// This is usually better than merely using `| head -n ...` on the output because the output
    /// is in random order while this option is biased toward picking origins that have shared
    /// history.
    ///
    /// The actual number of origins may be higher than this, in order to keep all origins with
    /// shared history with origins that were already selected.
    max_origins: Option<usize>,
    #[arg(long, default_value = "connected")]
    /// How to select origins
    origin_filter: OriginFilter,
    #[arg(long)]
    /// Path to a file where to write the URLs of all origins in the connected components
    origins_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 histhost_graph = Arc::new(Subgraph::with_node_constraint(
        Arc::clone(&graph),
        "ori,snp,rel,rev".parse().unwrap(),
    ));

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "node",
        local_speed = true,
    );
    pl.start("Listing origins reachable in connected components...");
    let seen = DashSet::new(); // shared between DFSs to avoid duplicate work
    let origins = DashSet::new();

    commits_in_vulnerabilities
        .into_par_iter()
        .for_each_with(pl.clone(), |pl, start_node| {
            if origins.len() > args.max_origins.unwrap_or(usize::MAX) {
                // We don't check this in the loop below in order to avoid keeping some origins
                // of a given project while excluding others.
                return;
            }
            let new = seen.insert(start_node);
            if !new {
                return;
            }
            let mut new_seen_nodes = 1;
            let mut todo = vec![start_node];

            while let Some(node) = todo.pop() {
                if swh_osv::packages::graph::should_ignore_node(&graph, node) {
                    continue;
                }
                if histhost_graph.properties().node_type(node) == NodeType::Origin {
                    origins.insert(node);
                }
                pl.light_update();
                for pred in histhost_graph.predecessors(node) {
                    let new = seen.insert(pred);
                    if new {
                        todo.push(pred);
                        new_seen_nodes += 1;
                    }
                }
                match args.origin_filter {
                    OriginFilter::Connected => {
                        for succ in histhost_graph.successors(node) {
                            let new = seen.insert(succ);
                            if new {
                                todo.push(succ);
                                new_seen_nodes += 1;
                            }
                        }
                    }
                    OriginFilter::Contains => (),
                }
            }

            log::debug!(
                "Saw {new_seen_nodes} new nodes from {}",
                graph.properties().swhid(start_node)
            );
        });
    pl.done();
    drop(pl);

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "origin",
        local_speed = true,
        expected_updates = Some(origins.len()),
    );
    pl.start("Writing origins...");
    let mut writer = BufWriter::new(
        File::create(&args.origins_out)
            .with_context(|| format!("Could not create {}", args.origins_out.display()))?,
    );
    for origin in origins {
        pl.light_update();
        let Some(origin_url) = graph.properties().message(origin) else {
            log::info!("Origin {origin} has no URL");
            continue;
        };
        let origin_url = String::from_utf8(origin_url.clone())
            .with_context(|| format!("Could not decode origin URL {origin_url:?}"))?;
        writer
            .write_all(format!("{}\n", origin_url).as_bytes())
            .with_context(|| format!("Could not write to {}", args.origins_out.display()))?;
    }

    writer
        .flush()
        .with_context(|| format!("Could not flush to {}", args.origins_out.display()))?;
    pl.done();

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

    Ok(())
}
