// 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

//! Version detection from package managers

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

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

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

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
/// Shows statistics about vulnerabilities and their package/git repo relationships
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 directory to write as ndjson files
    out_dir: PathBuf,
}

#[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_properties(|properties| properties.load_strings())
        .context("Could not load strings")?
        .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());

    // Create output writers
    let package_revisions_writer = {
        let dir = args.out_dir.join("package_revisions");
        dataset_writer::ParallelDatasetWriter::<dataset_writer::PlainZstTableWriter>::new(
            dir.clone(),
        )
        .with_context(|| format!("Could not create {}", dir.display()))?
    };

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

    let package_versions_not_found_writer = {
        let dir = args.out_dir.join("package_versions_not_found");
        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")),
            || {
                (
                    package_revisions_writer.get_thread_writer().unwrap(),
                    packages_not_found_writer.get_thread_writer().unwrap(),
                    package_versions_not_found_writer
                        .get_thread_writer()
                        .unwrap(),
                )
            },
            |(revisions_writer, not_found_writer, versions_not_found_writer),
             _rowid,
             filename,
             doc| {
                process_vulnerability(
                    graph,
                    revisions_writer.deref_mut(),
                    not_found_writer.deref_mut(),
                    versions_not_found_writer.deref_mut(),
                    filename,
                    doc,
                    &stats,
                )
            },
        )
    })?;

    pl.done();

    print_stats(&stats);

    Ok(())
}

#[instrument(skip(graph, revisions_writer, not_found_writer, versions_not_found_writer, doc, stats), fields(vuln = %filename))]
fn process_vulnerability<G, W: Write>(
    graph: &G,
    revisions_writer: &mut W,
    not_found_writer: &mut W,
    versions_not_found_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,
    <G as SwhGraphWithProperties>::Strings: properties::Strings,
{
    increment(&stats.total_vulns);

    if let Some(affected_list) = doc.affected {
        for affected_item in affected_list {
            // Extract package info and versions for error reporting
            let package_info = affected_item
                .package
                .as_ref()
                .map(|p| (p.ecosystem.clone(), p.name.clone()));

            let result = map_affected(graph, stats, &filename, affected_item)?;

            match result {
                MapAffectedResult::UnidentifiedPackage => (),
                MapAffectedResult::UnmappedEcosystem => (),
                MapAffectedResult::PackageNotInSwh(result) => {
                    if let Some((ecosystem, package_name)) = package_info {
                        let record = swh_osv::packages::PackageVersionMappingResult {
                            vulnerability_filename: filename.to_string(),
                            ecosystem: ecosystem.clone(),
                            package_name: package_name.clone(),
                            result,
                        };
                        serde_json::to_writer(&mut *not_found_writer, &record)
                            .with_context(|| format!("Could not write {record:?}"))?;
                        not_found_writer
                            .write_all(b"\n")
                            .context("Could not end line")?;
                    }
                }
                MapAffectedResult::Found(swh_affected) => {
                    // Process ranges and write records
                    let (ecosystem, package_name) =
                        &package_info.expect("found a mapping, but no package defined");
                    if let Some(ranges) = swh_affected.ranges {
                        for range in ranges {
                            for event in range.events {
                                match event.into_version() {
                                    swh_osv::packages::VersionInSwh::MappedFromVersion(result) => {
                                        // Write each VCS match to the revisions output
                                        let record =
                                            swh_osv::packages::PackageVersionMappingResult {
                                                vulnerability_filename: filename.to_string(),
                                                ecosystem: ecosystem.clone(),
                                                package_name: package_name.clone(),
                                                result: result.to_serializable(graph),
                                            };
                                        serde_json::to_writer(&mut *revisions_writer, &record)
                                            .with_context(|| {
                                                format!("Could not write VCS match {record:?}")
                                            })?;
                                        revisions_writer
                                            .write_all(b"\n")
                                            .context("Could not end line")?;
                                    }
                                    swh_osv::packages::VersionInSwh::VersionNotFound(result) => {
                                        // Write version not found record
                                        let record =
                                            swh_osv::packages::PackageVersionMappingResult {
                                                vulnerability_filename: filename.to_string(),
                                                ecosystem: ecosystem.clone(),
                                                package_name: package_name.clone(),
                                                result,
                                            };
                                        serde_json::to_writer(
                                            &mut *versions_not_found_writer,
                                            &record,
                                        )
                                        .with_context(
                                            || {
                                                format!(
                                                    "Could not write version not found {record:?}"
                                                )
                                            },
                                        )?;
                                        versions_not_found_writer
                                            .write_all(b"\n")
                                            .context("Could not end line")?;
                                    }
                                    _ => {
                                        // SinceTheBeginning, FromGit, PackageNotFound - nothing to write
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    Ok(())
}

fn print_stats(stats: &Stats) {
    let total = stats.total_vulns.load(Ordering::Relaxed);

    println!("\n=== Vulnerability Statistics ===\n");
    println!("Total vulnerabilities: {}", total);
    println!();

    let total_git_ranges = stats.total_git_ranges.load(Ordering::Relaxed);
    println!("--- Git Range Details ---");
    println!("Total git ranges: {}", total_git_ranges);
    if total_git_ranges > 0 {
        println!(
            "  Git ranges with repo field: {} ({:.1}%)",
            stats.git_ranges_with_repo.load(Ordering::Relaxed),
            percentage(
                stats.git_ranges_with_repo.load(Ordering::Relaxed),
                total_git_ranges
            )
        );
        println!(
            "  Git ranges without repo field: {} ({:.1}%)",
            stats.git_ranges_without_repo.load(Ordering::Relaxed),
            percentage(
                stats.git_ranges_without_repo.load(Ordering::Relaxed),
                total_git_ranges
            )
        );
        println!(
            "  Git ranges with commits: {} ({:.1}%)",
            stats.git_ranges_with_commits.load(Ordering::Relaxed),
            percentage(
                stats.git_ranges_with_commits.load(Ordering::Relaxed),
                total_git_ranges
            )
        );
        println!(
            "  Git ranges without commits: {} ({:.1}%)",
            stats.git_ranges_without_commits.load(Ordering::Relaxed),
            percentage(
                stats.git_ranges_without_commits.load(Ordering::Relaxed),
                total_git_ranges
            )
        );
    }
    println!();

    println!("--- Range Types ---");
    println!(
        "Total SEMVER ranges: {}",
        stats.total_semver_ranges.load(Ordering::Relaxed)
    );
    println!(
        "Total ECOSYSTEM ranges: {}",
        stats.total_ecosystem_ranges.load(Ordering::Relaxed)
    );
    println!("Total GIT ranges: {}", total_git_ranges);
    println!();

    let packages_with_info = stats.vulns_with_packages.load(Ordering::Relaxed);
    println!("--- Package Presence in SWH ---");
    if packages_with_info > 0 {
        // Compute totals from DashMaps
        let packages_found: u64 = stats
            .packages_found_in_swh_by_ecosystem
            .iter()
            .map(|entry| entry.value().load(Ordering::Relaxed))
            .sum();
        let packages_not_found: u64 = stats
            .packages_not_found_in_swh_by_ecosystem
            .iter()
            .map(|entry| entry.value().load(Ordering::Relaxed))
            .sum();
        let packages_unmapped = stats.packages_unmapped_ecosystem.load(Ordering::Relaxed);
        let total_checked = packages_found + packages_not_found + packages_unmapped;

        println!(
            "Packages found in SWH: {} ({:.1}%)",
            packages_found,
            percentage(packages_found, total_checked)
        );
        println!(
            "Packages not found in SWH: {} ({:.1}%)",
            packages_not_found,
            percentage(packages_not_found, total_checked)
        );
        println!(
            "Packages with unmapped ecosystem: {} ({:.1}%)",
            packages_unmapped,
            percentage(packages_unmapped, total_checked)
        );

        // Print per-ecosystem breakdown
        println!("\n  By ecosystem:");
        let mut ecosystems: Vec<String> = stats
            .packages_found_in_swh_by_ecosystem
            .iter()
            .map(|entry| entry.key().clone())
            .chain(
                stats
                    .packages_not_found_in_swh_by_ecosystem
                    .iter()
                    .map(|entry| entry.key().clone()),
            )
            .collect();
        ecosystems.sort();
        ecosystems.dedup();

        // Collect data for table
        let mut rows = Vec::new();
        for ecosystem in ecosystems {
            let found = stats
                .packages_found_in_swh_by_ecosystem
                .get(&ecosystem)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0);
            let not_found = stats
                .packages_not_found_in_swh_by_ecosystem
                .get(&ecosystem)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0);
            let eco_total = found + not_found;

            if eco_total > 0 {
                rows.push((ecosystem, found, not_found, eco_total));
            }
        }

        if !rows.is_empty() {
            // Print markdown table
            println!(
                "  | {:<100} | {:>15} | {:>15} | {:>15} |",
                "Ecosystem", "Found", "Not Found", "Total"
            );
            println!(
                "  |-{:-<100}-|-{:->15}-|-{:->15}-|-{:->15}-|",
                "", "", "", ""
            );
            for (ecosystem, found, not_found, eco_total) in rows {
                println!(
                    "  | {:<100} | {:>15} | {:>15} | {:>15} |",
                    ecosystem,
                    format!("{} ({:.1}%)", found, percentage(found, eco_total)),
                    format!("{} ({:.1}%)", not_found, percentage(not_found, eco_total)),
                    eco_total
                );
            }
        }
    } else {
        println!("No packages to check in SWH");
    }
    println!();

    println!("--- Package Version Resolution ---");
    // Compute totals from DashMaps
    let found_sure: u64 = stats
        .found_sure_package_version_in_swh_by_ecosystem
        .iter()
        .map(|entry| entry.value().load(Ordering::Relaxed))
        .sum();
    let found_fuzzy: u64 = stats
        .found_fuzzy_package_version_in_swh_by_ecosystem
        .iter()
        .map(|entry| entry.value().load(Ordering::Relaxed))
        .sum();
    let not_found: u64 = stats
        .unfound_package_version_in_swh_by_ecosystem
        .iter()
        .map(|entry| entry.value().load(Ordering::Relaxed))
        .sum();
    let total_version_lookups = found_sure + found_fuzzy + not_found;

    if total_version_lookups > 0 {
        println!(
            "Package versions found in SWH: {} ({:.1}%) for sure, {} ({:.1}%) fuzzily",
            found_sure,
            percentage(found_sure, total_version_lookups),
            found_fuzzy,
            percentage(found_fuzzy, total_version_lookups)
        );
        println!(
            "Package versions not found in SWH: {} ({:.1}%)",
            not_found,
            percentage(not_found, total_version_lookups)
        );

        // Print per-ecosystem breakdown
        println!("\n  By ecosystem:");
        let mut ecosystems: Vec<String> = stats
            .found_sure_package_version_in_swh_by_ecosystem
            .iter()
            .map(|entry| entry.key().clone())
            .chain(
                stats
                    .found_fuzzy_package_version_in_swh_by_ecosystem
                    .iter()
                    .map(|entry| entry.key().clone()),
            )
            .chain(
                stats
                    .unfound_package_version_in_swh_by_ecosystem
                    .iter()
                    .map(|entry| entry.key().clone()),
            )
            .collect();
        ecosystems.sort();
        ecosystems.dedup();

        // Collect data for table
        let mut rows = Vec::new();
        for ecosystem in ecosystems {
            let sure = stats
                .found_sure_package_version_in_swh_by_ecosystem
                .get(&ecosystem)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0);
            let fuzzy = stats
                .found_fuzzy_package_version_in_swh_by_ecosystem
                .get(&ecosystem)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0);
            let unfound = stats
                .unfound_package_version_in_swh_by_ecosystem
                .get(&ecosystem)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0);
            let eco_total = sure + fuzzy + unfound;

            if eco_total > 0 {
                rows.push((ecosystem, sure, fuzzy, unfound, eco_total));
            }
        }

        if !rows.is_empty() {
            // Print markdown table
            println!(
                "  | {:<100} | {:>15} | {:>15} | {:>15} | {:>10} |",
                "Ecosystem", "Sure Match", "Fuzzy Match", "Not Found", "Total"
            );
            println!(
                "  |-{:-<100}-|-{:->15}-|-{:->15}-|-{:->15}-|-{:->10}-|",
                "", "", "", "", ""
            );
            for (ecosystem, sure, fuzzy, unfound, eco_total) in rows {
                println!(
                    "  | {:<100} | {:>15} | {:>15} | {:>15} | {:>10} |",
                    ecosystem,
                    format!("{} ({:.1}%)", sure, percentage(sure, eco_total)),
                    format!("{} ({:.1}%)", fuzzy, percentage(fuzzy, eco_total)),
                    format!("{} ({:.1}%)", unfound, percentage(unfound, eco_total)),
                    eco_total
                );
            }
        }
    } else {
        println!("No package version lookups performed");
    }
    println!();

    println!("--- Package Version Directories and Revisions ---");
    let with_revs: u64 = stats
        .package_dirs_with_revisions
        .iter()
        .map(|entry| entry.value().load(Ordering::Relaxed))
        .sum();
    let without_revs: u64 = stats
        .package_dirs_without_revisions
        .iter()
        .map(|entry| entry.value().load(Ordering::Relaxed))
        .sum();
    let total_dirs = with_revs + without_revs;
    if total_dirs > 0 {
        println!(
            "Package directories with revisions pointing to them: {} ({:.1}%)",
            with_revs,
            percentage(with_revs, total_dirs)
        );
        println!(
            "Package directories without revisions: {} ({:.1}%)",
            without_revs,
            percentage(without_revs, total_dirs)
        );
        println!(
            "Total revisions pointing to package directories: {}",
            stats
                .total_revisions_pointing_to_package_dirs
                .iter()
                .map(|entry| entry.value().load(Ordering::Relaxed))
                .sum::<u64>()
        );

        let in_vcs: u64 = stats
            .package_version_dir_in_vcs_by_ecosystem
            .iter()
            .map(|entry| entry.value().load(Ordering::Relaxed))
            .sum();
        let not_in_vcs: u64 = stats
            .package_version_dir_not_in_vcs_by_ecosystem
            .iter()
            .map(|entry| entry.value().load(Ordering::Relaxed))
            .sum();
        let total_revs = in_vcs + not_in_vcs;

        if total_revs > 0 {
            println!(
                "  Revisions with parent (in VCS): {} ({:.1}%)",
                in_vcs,
                percentage(in_vcs, total_revs)
            );
            println!(
                "  Revisions without parent (standalone): {} ({:.1}%)",
                not_in_vcs,
                percentage(not_in_vcs, total_revs)
            );

            // Print per-ecosystem breakdown
            println!("\n  By ecosystem:");
            let mut ecosystems: Vec<String> = stats
                .package_version_dir_in_vcs_by_ecosystem
                .iter()
                .map(|entry| entry.key().clone())
                .chain(
                    stats
                        .package_version_dir_not_in_vcs_by_ecosystem
                        .iter()
                        .map(|entry| entry.key().clone()),
                )
                .collect();
            ecosystems.sort();
            ecosystems.dedup();

            // Collect data for table
            let mut rows = Vec::new();
            for ecosystem in ecosystems {
                let with_parent = stats
                    .package_version_dir_in_vcs_by_ecosystem
                    .get(&ecosystem)
                    .map(|v| v.load(Ordering::Relaxed))
                    .unwrap_or(0);
                let without_parent = stats
                    .package_version_dir_not_in_vcs_by_ecosystem
                    .get(&ecosystem)
                    .map(|v| v.load(Ordering::Relaxed))
                    .unwrap_or(0);
                let without_rev = stats
                    .package_dirs_without_revisions
                    .get(&ecosystem)
                    .map(|v| v.load(Ordering::Relaxed))
                    .unwrap_or(0);
                let eco_total = with_parent + without_parent + without_rev;

                if eco_total > 0 {
                    rows.push((
                        ecosystem,
                        with_parent,
                        without_parent,
                        without_rev,
                        eco_total,
                    ));
                }
            }

            if !rows.is_empty() {
                // Print markdown table
                println!(
                    "  | {:<100} | {:>19} | {:>19} | {:>19} | {:>10} |",
                    "Ecosystem", "Rev With Parent", "Rev Without Parent", "Without Rev", "Total"
                );
                println!(
                    "  |-{:-<100}-|-{:->19}-|-{:->19}-|-{:->19}-|-{:->10}-|",
                    "", "", "", "", ""
                );
                for (ecosystem, with_parent, without_parent, without_rev, eco_total) in rows {
                    println!(
                        "  | {:<100} | {:>19} | {:>19} | {:>19} | {:>10} |",
                        ecosystem,
                        format!(
                            "{} ({:.1}%)",
                            with_parent,
                            percentage(with_parent, eco_total)
                        ),
                        format!(
                            "{} ({:.1}%)",
                            without_parent,
                            percentage(without_parent, eco_total)
                        ),
                        format!(
                            "{} ({:.1}%)",
                            without_rev,
                            percentage(without_rev, eco_total)
                        ),
                        eco_total
                    );
                }
            }
        }
    } else {
        println!("No package directories found");
    }
}

fn percentage(value: u64, total: u64) -> f64 {
    if total == 0 {
        0.0
    } else {
        (value as f64 / total as f64) * 100.0
    }
}
