// Copyright (C) 2025  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::borrow::Borrow;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};

use anyhow::{Context, Result, anyhow};
use dashmap::DashSet;
use dsi_progress_logger::ConcurrentProgressLog;
use dsi_progress_logger::{ProgressLog, progress_logger};
use osv::schema::{Range, RangeType, Vulnerability};
use rapidhash::{RapidBuildHasher, RapidHashSet};
use rusqlite::Connection;
use swh_graph::graph::{NodeId, SwhGraphWithProperties};
use swh_graph::properties;
use swh_graph::views::contiguous_subgraph::ContractionBackend;

use crate::VulnerabilityId;
use crate::git;
use crate::multimap::NodeMultimap;

#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct VulnerabilityRowId(pub u64);

pub fn count_vulnerabilities(conn: &Connection) -> Result<usize> {
    conn.prepare("SELECT COUNT(*) FROM vulnerabilities")
        .context("Could not access 'vulnerabilities' table")?
        .query_one([], |row| {
            let count: usize = row.get(0).expect("Missing count column");
            Ok(count)
        })
        .context("Could not count vulnerabilities")
}

pub fn for_each_vulnerability<
    TL,
    F: Fn(&mut TL, VulnerabilityRowId, String, Vulnerability) -> Result<()> + Sync,
>(
    conn: &Connection,
    pl: &impl ConcurrentProgressLog,
    parallelism: Option<NonZeroUsize>,
    init_threadlocal: impl Fn() -> TL + Sync,
    f: F,
) -> Result<()> {
    let parallelism = parallelism.map(usize::from).unwrap_or(num_cpus::get());
    let (rows_tx, rows_rx) =
        crossbeam::channel::bounded::<(VulnerabilityRowId, String, String)>(10 * parallelism);
    let f = &f;
    let errored = AtomicBool::new(false); // to abort early
    std::thread::scope(|scope| -> Result<_> {
        let worker_threads: Vec<_> = (0..parallelism)
            .map(|_| {
                let mut pl = pl.clone();
                let rows_rx = rows_rx.clone();
                let init_threadlocal = &init_threadlocal;
                let errored = &errored;
                scope.spawn(move || -> Result<_> {
                    let mut thread_local = init_threadlocal();
                    while let Ok((rowid, filename, doc)) = rows_rx.recv() {
                        let vuln: Vulnerability = match serde_json::from_str(&doc) {
                            Err(e) => {
                                if e.classify() == serde_json::error::Category::Data {
                                    if filename.starts_with("openSUSE-SU-") {
                                        // noisy because of all the missing URLs in
                                        // the 'references' field
                                        log::debug!("Skipping unparsable file {filename}: {e}");
                                    } else {
                                        log::warn!("Skipping unparsable file {filename}: {e}");
                                    }
                                    continue;
                                } else {
                                    errored.store(true, Ordering::Relaxed);
                                    Err(e)?
                                }
                            }
                            Ok(vuln) => vuln,
                        };
                        let _span = tracing::info_span!("vuln", filename);
                        f(&mut thread_local, rowid, filename, vuln)
                            .inspect_err(|_| errored.store(true, Ordering::Relaxed))?;
                        pl.light_update();
                    }
                    Ok(())
                })
            })
            .collect();

        drop(rows_rx); // makes rows_tx.send() fail if all worker threads crash

        let mut res = Ok(());
        for row in conn
            .prepare("SELECT ROWID, filename, doc FROM vulnerabilities")
            .context("Unexpected schema of 'vulnerabilities' table")?
            .query_map([], |row: &rusqlite::Row| {
                let rowid: u64 = row.get(0).expect("Missing 'ROWID' column");
                let filename: String = row.get(1).expect("Missing 'filename' column");
                let doc: String = row.get(2).expect("Missing 'doc' column");
                Ok((VulnerabilityRowId(rowid), filename, doc))
            })?
        {
            if errored.load(Ordering::Relaxed) {
                // One of the worker threads exited with an error.
                break;
            }
            if rows_tx.send(row?).is_err() {
                res = Err(anyhow!("Receivers unexpected terminated"));
                break;
            }
        }

        drop(rows_tx); // tell workers to stop

        for handle in worker_threads {
            handle.join().expect("Could not join thread")?;
        }
        res
    })
}

pub fn for_each_git_range<
    TL,
    F: Fn(&mut TL, VulnerabilityRowId, &str, &Vulnerability, &Range) -> Result<()> + Sync,
>(
    conn: &Connection,
    pl: &impl ConcurrentProgressLog,
    parallelism: Option<NonZeroUsize>,
    init_threadlocal: impl Fn() -> TL + Sync,
    f: F,
) -> Result<()> {
    for_each_vulnerability(
        conn,
        pl,
        parallelism,
        init_threadlocal,
        |thread_local, rowid, filename, doc| {
            if let Some(affected) = &doc.affected {
                for affected in affected {
                    if let Some(ranges) = &affected.ranges {
                        for range in ranges {
                            if let RangeType::Git = range.range_type {
                                f(thread_local, rowid, &filename, &doc, range)?;
                            }
                        }
                    }
                }
            }
            Ok(())
        },
    )
}

/// Collects all commits mentioned in vulnerability git ranges into a set.
pub fn collect_commits_from_vulnerability_ranges<G, CB, NMM, BorrowedNNM>(
    conn: &Connection,
    pl: &impl ConcurrentProgressLog,
    parallelism: Option<NonZeroUsize>,
    git_range_parser: &git::GitRangeParser<G, CB, NMM, BorrowedNNM>,
    commits_set: &DashSet<NodeId, RapidBuildHasher>,
) -> Result<()>
where
    G: SwhGraphWithProperties<Maps: properties::Maps> + Sync,
    CB: ContractionBackend + Sync,
    NMM: NodeMultimap<NodeId> + Sync,
    BorrowedNNM: Borrow<NMM> + Sync,
{
    for_each_git_range(
        conn,
        pl,
        parallelism,
        || (),
        |(), _rowid, filename, _vulnerability, range| {
            let Some((commits_by_event_type, _has_zero_commit)) = git_range_parser
                .parse_git_range(range, None)
                .with_context(|| format!("Could not parse git range in {filename}"))?
            else {
                // some unknown commits in the range, ignore.
                return Ok(());
            };
            let commits_in_vuln: RapidHashSet<_> = commits_by_event_type
                .iter()
                .flat_map(|set| set.iter())
                .collect();

            // insert those early to avoid work on the same connected component from other threads
            for commit in commits_in_vuln {
                commits_set.insert(commit);
            }
            Ok(())
        },
    )
}
pub fn read_filenames_from_sqlite(path: &Path) -> Result<HashMap<String, VulnerabilityId>> {
    log::debug!("Opening sqlite database...");
    let conn =
        Connection::open(path).with_context(|| format!("Could not open {}", path.display()))?;

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

    let mut pl = progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(num_vulnerabilities),
    );
    pl.start(format!("Reading {num_vulnerabilities} vulnerabilities..."));
    let filenames = conn
        .prepare("SELECT id, filename FROM vulnerabilities")
        .context("Unexpected schema of 'vulnerabilities' table")?
        .query_map([], |row: &rusqlite::Row| {
            let id: u64 = row.get(0).expect("Missing 'id' column");
            let filename: String = row.get(1).expect("Missing 'filename' column");
            pl.light_update();
            Ok((filename, VulnerabilityId(id)))
        })?
        .collect::<rusqlite::Result<_>>()?;
    Ok(filenames)
}
