// Copyright (C) 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

//! List all versions affected by vulnerabilities

use std::io::Write;
use std::num::NonZeroUsize;
use std::ops::DerefMut;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::Parser;
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use mimalloc::MiMalloc;
use osv::schema::{Ecosystem, Vulnerability};
use rusqlite::Connection;
use serde::Serialize;
use swh_graph::SWHID;
use swh_graph::graph::*;
use swh_graph::mph::DynMphf;
use swh_graph::properties;
use tracing::{Level, span};

use swh_osv::db;
use swh_osv::packages::{MapAffectedResult, Stats, list_affected_versions, map_affected};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
/// Lists all versions affected by vulnerabilities in the OSV 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 the output file (ndjson format)
    out_dir: PathBuf,
}

#[derive(Debug, Serialize)]
struct AffectedVersionRecord {
    vulnerability_filename: String,
    ecosystem: Ecosystem,
    package_name: String,
    versions: Vec<VersionInfo>,
}

#[derive(Debug, Serialize)]
struct VersionInfo {
    original_version: Option<String>,
    branch_name: Option<String>,
    swhid: SWHID,
}

#[tokio::main(flavor = "multi_thread")]
pub async 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_label_names())
        .context("Could not load label names")?
        .load_labels()
        .context("Could not load labels")?;

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

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

    let stats = Arc::new(Stats::default());

    let writer = {
        let dir = args.out_dir.clone();
        dataset_writer::ParallelDatasetWriter::<dataset_writer::PlainZstTableWriter>::new(
            dir.clone(),
        )
        .with_context(|| format!("Could not create {}", dir.display()))?
    };

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

    let graph = &graph;
    tokio::task::block_in_place(|| {
        db::for_each_vulnerability(
            &conn,
            &pl,
            Some(NonZeroUsize::try_from(num_cpus::get()).expect("0 CPUs")),
            || writer.get_thread_writer().unwrap(),
            |writer, _rowid, filename, doc| {
                process_vulnerability(graph, writer.deref_mut(), filename, doc, &stats)
            },
        )
    })?;

    pl.done();

    tracing::info!("Done processing vulnerabilities");

    println!("stats: {stats:#?}");

    Ok(())
}

fn process_vulnerability<G, W: Write>(
    graph: &G,
    writer: &mut W,
    filename: String,
    doc: Vulnerability,
    stats: &Stats,
) -> Result<()>
where
    G: SwhLabeledForwardGraph + SwhBackwardGraph + SwhGraphWithProperties,
    <G as SwhGraphWithProperties>::Maps: properties::Maps,
    <G as SwhGraphWithProperties>::LabelNames: properties::LabelNames,
{
    let _span = span!(Level::INFO, "vuln", filename = filename).entered();

    if let Some(affected_list) = doc.affected {
        for affected_item in affected_list {
            let package_info = affected_item
                .package
                .as_ref()
                .map(|p| (p.ecosystem.clone(), p.name.clone()));

            let result = {
                let _span = span!(
                    Level::INFO,
                    "affected",
                    package = %format!("{:?}", affected_item.package),
                )
                .entered();
                map_affected(graph, stats, &filename, affected_item)?
            };

            if let MapAffectedResult::Found(swh_affected) = result {
                let (ecosystem, package_name) =
                    &package_info.expect("found a mapping, but no package defined");

                let _span = span!(
                    Level::INFO,
                    "affected",
                    ecosystem = %format!("{ecosystem:?}"),
                    origins = %format!("{:?}", swh_affected.origins),
                )
                .entered();

                let affected_versions =
                    list_affected_versions(graph, stats, &filename, ecosystem, swh_affected)
                        .with_context(|| {
                            format!("Could not list versions affected by {filename}")
                        })?;

                if !affected_versions.is_empty() {
                    let versions = affected_versions
                        .into_iter()
                        .map(|v| VersionInfo {
                            original_version: v.original_version,
                            branch_name: v.branch_name,
                            swhid: v.swhid,
                        })
                        .collect();

                    let record = AffectedVersionRecord {
                        vulnerability_filename: filename.clone(),
                        ecosystem: ecosystem.clone(),
                        package_name: package_name.clone(),
                        versions,
                    };
                    serde_json::to_writer(&mut *writer, &record)
                        .with_context(|| format!("Could not write {record:?}"))?;
                    writer.write_all(b"\n").context("Could not end line")?;
                }
            }
        }
    }
    Ok(())
}
