// 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::{BufReader, 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 swh_graph::graph::{
    SwhBackwardGraph, SwhBidirectionalGraph, SwhForwardGraph, SwhGraphWithProperties,
};
use swh_graph::mph::DynMphf;
use swh_graph::views::Subgraph;
use swh_graph::{NodeType, SWHID};

use swh_osv::kb::{Artifact, Fix, RepositoryCommit};
use swh_osv::packages::map_purl_to_origin_urls;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[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 a local clone of the 'statements' directory of https://github.com/SAP/project-kb's
    /// 'vulnerability-data' branch
    kb: PathBuf,
    #[arg(long)]
    /// Number of worker threads. Defaults to the number of CPU threads.
    parallelism: Option<NonZeroUsize>,
    #[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::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);

    let subdirs = std::fs::read_dir(&args.kb)
        .with_context(|| format!("Could not list {}", args.kb.display()))?
        .map(|entry| {
            let path = entry
                .with_context(|| format!("Could not stat {} entry", args.kb.display()))?
                .path();
            if path
                .file_name()
                .expect("file name")
                .to_string_lossy()
                .starts_with('.')
            {
                Ok(None)
            } else {
                Ok(Some(path.join("statement.yaml")))
            }
        })
        .collect::<Result<Vec<Option<PathBuf>>>>()?
        .into_iter()
        .flatten()
        .collect::<Vec<_>>();

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "statement",
        local_speed = true,
        expected_updates = Some(subdirs.len()),
    );
    pl.threshold(1000);
    pl.start("Reading KB statements");

    // set of all commits mentioned by the vulnerability database
    let referenced_commits = DashSet::with_hasher(RapidBuildHasher::default());
    let referenced_origins = DashSet::new();

    /*
    let commit_typos = HashMap::from([
        ("5a711c05ec750f069235597173084c2ee796242", "5a711c05ec750f069235597173084c2ee7962424"),
        ("2c6964ae94d8f9a9c9a32e5ae5a0b794e8b8d3b", "2c6964ae94d8f9a9c9a32e5ae5a0b794e8b8d3be"),
        ("87c92b7b38890c217bc76f2c55036e6a5cca9a0", "87c92b7b38890c217bc76f2c55036e6a5cca9a07"),
        ("9f7376abbff7434794f2c7c2909e02bac232fb5", "9f7376abbff7434794f2c7c2909e02bac232fb5b"),
        ("ec3d0db81ba061b27e934d5ff56e9baca0049eb", "ec3d0db81ba061b27e934d5ff56e9baca0049ebf"),
        ("8afc5d1757795fde715902067360af5d90f046d", "8afc5d1757795fde715902067360af5d90f046da"),
    ]);
    */

    subdirs
        .into_par_iter()
        .try_for_each_with(pl.clone(), |pl, subdir| -> Result<_> {
            let statement_path = subdir;
            let statement_file = std::fs::File::open(&statement_path)
                .with_context(|| format!("Could not open {}", statement_path.display()))?;
            let statement: swh_osv::kb::Statement =
                serde_yaml_ng::from_reader(BufReader::new(statement_file))
                    .with_context(|| format!("Could not parse {}", statement_path.display()))?;

            for artifact in statement.artifacts.unwrap_or_default() {
                let Artifact {
                    id,
                    reason: _,
                    affected: _,
                } = artifact;
                for origin_url in map_purl_to_origin_urls(id) {
                    if let Ok(node) = graph
                        .properties()
                        .node_id(SWHID::from_origin_url(origin_url))
                    {
                        referenced_origins.insert(node);
                    }
                }
            }

            for fix in statement.fixes.unwrap_or_default() {
                let Fix { id: _, commits } = fix;
                for commit in commits {
                    let RepositoryCommit { id, repository } = commit;
                    //let id = commit_typos.get(&id.as_str()).cloned().unwrap_or(id.as_str());
                    let swhid: SWHID = match format!("swh:1:rev:{id}").parse() {
                        Ok(swhid) => swhid,
                        Err(e) => {
                            log::warn!(
                                "Invalid commit id in {}: {id} ({e})",
                                statement_path.display()
                            );
                            continue;
                        }
                    };
                    if let Ok(commit) = graph.properties().node_id(swhid) {
                        referenced_commits.insert(commit);
                    }
                    if let Ok(node) = graph
                        .properties()
                        .node_id(SWHID::from_origin_url(repository))
                    {
                        referenced_origins.insert(node);
                    }
                }
            }

            pl.light_update();
            Ok(())
        })?;

    pl.done();
    drop(pl);

    log::info!(
        "Found {} commits and {} origins",
        referenced_commits.len(),
        referenced_origins.len()
    );

    let mut referenced_nodes = referenced_commits;
    referenced_nodes.extend(referenced_origins);

    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,
        expected_updates = Some(referenced_nodes.len()),
    );
    pl.start("Listing origins reachable in connected components...");

    let seen = DashSet::new(); // shared between DFSs to avoid duplicate work
    let origins = DashSet::new();

    referenced_nodes
        .into_par_iter()
        .for_each_with(pl.clone(), |pl, start_node| {
            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);
                }
                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 => (),
                }
            }

            pl.light_update();

            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(())
}
