// 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::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};

use anyhow::Result;
use dashmap::DashMap;
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use rapidhash::{RapidBuildHasher, RapidHashMap, RapidHashSet};
use rayon::prelude::*;
use swh_graph::graph::{SwhBackwardGraph, SwhForwardGraph, SwhGraphWithProperties};
use swh_graph::properties;
use swh_graph::{NodeType, SWHID};
use swh_graph_stdlib::connectivity::SubgraphWccs;
use value_traits::slices::SliceByValue;

type NodeId = usize;

#[derive(Default, Debug)]
pub struct Stats {
    pub no_matches: AtomicU64,
    pub no_matches_but_mentions_cherrypick: AtomicU64,
    pub at_least_one_match: AtomicU64,
    pub at_least_two_matches: AtomicU64,
    pub unparsable_commit: AtomicU64,
    pub unknown_commit_id_and_no_repo_url: AtomicU64,
    pub unknown_commit_id_and_has_nonutf8_origin_url: AtomicU64,
    pub unknown_commit_id_and_has_known_origin_url: AtomicU64,
    pub unknown_commit_id_and_has_unknown_origin_url: AtomicU64,
    pub different_component: AtomicU64,
    pub found: AtomicU64,
}

pub struct Original2Cherrypicks(pub RapidHashMap<NodeId, Vec<NodeId>>);
pub struct Cherrypick2Originals(pub RapidHashMap<NodeId, Vec<NodeId>>);

/// Returns the list of cherry-picks of each node mentioned in a vulnerability, if any
pub fn list_cherrypicks<
    G: SwhForwardGraph
        + SwhBackwardGraph
        + SwhGraphWithProperties<Maps: properties::Maps, Strings: properties::Strings>
        + Sync
        + Send,
    D: AsRef<[usize]> + Sync + Send,
    V: SliceByValue<Value = usize> + Sync + Send,
>(
    graph: &G,
    wccs: &SubgraphWccs<D, V>,
) -> Result<(Stats, Cherrypick2Originals)> {
    let cherrypicks_to_originals = DashMap::with_hasher(RapidBuildHasher::default());

    let num_connected_nodes = wccs.num_nodes();

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(num_connected_nodes),
    );
    pl.threshold(1000);
    pl.start(format!(
        "Scanning for cherry-pick among {} connected nodes",
        num_connected_nodes
    ));

    let stats = Stats::default();

    let cherrypick_re =
        regex::bytes::Regex::new(r"[\r\n]\(cherry picked from commit (?<commit_sha1>[0-9a-f]{40})(\n? +(?<repo_url>https?|git)://\S+)?(\n? +(?<branch_name>\S+))?\)[\r\n]")
            .expect("Invalid regexp");

    wccs
        .par_iter_nodes()
        .for_each_with(pl.clone(), |pl, node| {
            pl.light_update();
            let Some(message) = &graph.properties().message(node) else {
                log::debug!("{} has no message", graph.properties().swhid(node));
                return;
            };
            let mut num_matches = 0;
            for capture in cherrypick_re.captures_iter(message) {
                num_matches += 1;
                let commit_id = &capture["commit_sha1"];
                let mut swhid_hash = [0u8; 20];
                if let Err(e) = faster_hex::hex_decode(commit_id, &mut swhid_hash) {
                    stats.unparsable_commit.fetch_add(1, Ordering::Relaxed);
                    log::warn!(
                        "Could not parse commit id {} referenced by {}: {e}",
                        String::from_utf8_lossy(commit_id),
                        graph.properties().swhid(node)
                    );
                    continue;
                }
                let swhid = SWHID {
                    namespace_version: 1,
                    node_type: NodeType::Revision,
                    hash: swhid_hash,
                };
                let Ok(cherrypicked_from) = graph.properties().node_id(swhid) else {
                    if let Some(repo_url_match) = capture.name("repo_url") {
                        let Ok(repo_url) = String::from_utf8(repo_url_match.as_bytes().to_vec()) else {
                            stats.unknown_commit_id_and_has_nonutf8_origin_url.fetch_add(1, Ordering::Relaxed);
                            log::warn!("{} references non-UTF8 origin URL: {}", String::from_utf8_lossy(commit_id), String::from_utf8_lossy(repo_url_match.as_bytes()));
                            continue;
                        };

                        let origin_swhid = SWHID::from_origin_url(&repo_url);
                        if graph.properties().node_id(origin_swhid).is_ok() {
                            stats.unknown_commit_id_and_has_known_origin_url.fetch_add(1, Ordering::Relaxed);
                            log::debug!(
                                "Could not find {} referenced by {} in the graph (origin {repo_url} may be outdated)",
                                String::from_utf8_lossy(commit_id),
                                graph.properties().swhid(node),
                            );
                        } else {
                            stats.unknown_commit_id_and_has_unknown_origin_url.fetch_add(1, Ordering::Relaxed);
                            log::debug!(
                                "Could not find {} referenced by {} in the graph (origin {repo_url} is not in the graph)",
                                String::from_utf8_lossy(commit_id),
                                graph.properties().swhid(node),
                            );
                        }
                    } else {
                        stats.unknown_commit_id_and_no_repo_url.fetch_add(1, Ordering::Relaxed);
                        log::debug!(
                            "Could not find {} referenced by {} in the graph",
                            String::from_utf8_lossy(commit_id),
                            graph.properties().swhid(node)
                        );
                    }
                    continue;
                };

                if wccs.component(node) != wccs.component(cherrypicked_from) {
                    stats.different_component.fetch_add(1, Ordering::Relaxed);
                    log::debug!("{} claims to be cherry-picked from {} but is not in the same connected component; ignoring",
                               graph.properties().swhid(node), graph.properties().swhid(cherrypicked_from));
                    continue;
                }
                if wccs.component(node).is_none() {
                    log::warn!("{} is cherry-picked from {}, but neither are in the connected components",
                               graph.properties().swhid(node), graph.properties().swhid(cherrypicked_from));
                    continue;
                }

                stats.found.fetch_add(1, Ordering::Relaxed);
                cherrypicks_to_originals
                    .entry(node)
                    .or_insert_with(|| Mutex::new(RapidHashSet::with_hasher(Default::default())))
                    .lock()
                    .unwrap()
                    .insert(cherrypicked_from);
            }

            if num_matches == 0 {
                stats.no_matches.fetch_add(1, Ordering::Relaxed);
                if String::from_utf8_lossy(message).contains("cherry picked from commit") {
                    stats.no_matches_but_mentions_cherrypick.fetch_add(1, Ordering::Relaxed);
                    log::debug!("Could not parse cherry-pick stanza of {}. Commit message: {}", graph.properties().swhid(node), String::from_utf8_lossy(message));
                }
            } else {
                stats.at_least_one_match.fetch_add(1, Ordering::Relaxed);
                if num_matches >= 2 {
                    log::debug!("{} has {num_matches} valid cherry pick stanzas. Commit message: {}", graph.properties().swhid(node), String::from_utf8_lossy(message));
                    stats.at_least_two_matches.fetch_add(1, Ordering::Relaxed);
                }
            }
        });

    pl.done();

    Ok((
        stats,
        Cherrypick2Originals(
            cherrypicks_to_originals
                .into_par_iter()
                // strip mutexes and convert HashSet to Vec
                .map(|(cherrypick, cherrypicked)| -> (NodeId, Vec<NodeId>) {
                    (
                        cherrypick,
                        cherrypicked.into_inner().unwrap().into_iter().collect(),
                    )
                })
                .collect(),
        ),
    ))
}

pub fn invert_cherrypick_map(
    num_cherrypicks: usize,
    cherrypicks_to_originals: impl IntoParallelIterator<
        Item = (impl Borrow<NodeId>, impl IntoIterator<Item: Borrow<NodeId>>),
    >,
) -> impl Iterator<Item = (NodeId, Vec<NodeId>)> {
    let cherrypicks_to_originals = cherrypicks_to_originals.into_par_iter();
    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "cherry-pick",
        local_speed = true,
        expected_updates = Some(num_cherrypicks),
    );
    pl.start("Inverting cherrypick-to-original multimap");
    let original_to_cherrypicks = DashMap::with_hasher(RapidBuildHasher::default());
    cherrypicks_to_originals.for_each_with(pl.clone(), |pl, (cherrypick, originals)| {
        for original in originals {
            original_to_cherrypicks
                .entry(*original.borrow())
                .or_insert_with(|| Mutex::new(Vec::new()))
                .lock()
                .unwrap()
                .push(*cherrypick.borrow());
        }
        pl.light_update();
    });
    pl.done();
    original_to_cherrypicks
        .into_iter()
        // strip mutexes
        .map(|(original, cherrypicks)| (original, cherrypicks.into_inner().unwrap()))
}

#[derive(Debug, Default, PartialEq, Eq)]
pub struct ReductionStats {
    pub total_cherrypicks: usize,
    pub total_nonsingleton_cherrypicks: usize,
    pub reducible_cherrypickedfrom: usize,
    pub partially_reducible_cherrypickedfrom: usize,
}

impl std::ops::Add for ReductionStats {
    type Output = Self;

    fn add(mut self, other: Self) -> Self {
        self.total_cherrypicks += other.total_cherrypicks;
        self.total_nonsingleton_cherrypicks += other.total_nonsingleton_cherrypicks;
        self.reducible_cherrypickedfrom += other.reducible_cherrypickedfrom;
        self.partially_reducible_cherrypickedfrom += other.partially_reducible_cherrypickedfrom;
        self
    }
}

/// Reduces transitive cherry-pick relations.
///
/// For example, if A is cherry-picked from B and C, and B is itself cherry-picked from C,
/// then this function can reduce A's cherry-pick-from set to just {B}, since the relation
/// to C is implied transitively through B.
pub fn reduce_transitive_cherrypicks<G>(
    graph: &G,
    cherrypicks_to_originals: &RapidHashMap<NodeId, Vec<NodeId>>,
) -> (ReductionStats, RapidHashMap<NodeId, Vec<NodeId>>)
where
    G: SwhGraphWithProperties<Maps: properties::Maps> + Sync + Send,
{
    let (stats, reduced_cherrypicks_to_originals) = cherrypicks_to_originals
        .par_iter()
        .fold(
            || (ReductionStats::default(), Vec::new()),
            |(mut stats, mut reduced_cherrypicks_to_originals),
             (&node1, cherrypicked_from): (&NodeId, &Vec<NodeId>)| {
                let mut reduced_cherrypicked_from: Vec<usize> = cherrypicked_from.clone();
                let cherrypicked_from: RapidHashSet<_> = cherrypicked_from.iter().collect();
                stats.total_cherrypicks += 1;
                if cherrypicked_from.len() > 1 {
                    stats.total_nonsingleton_cherrypicks += 1;
                    for &&node2 in &cherrypicked_from {
                        if let Some(transitive_cherrypicked_from) = cherrypicks_to_originals
                            .get(&node2)
                            .map(|set| -> RapidHashSet<_> { set.iter().collect() })
                        {
                            let difference: Vec<usize> = cherrypicked_from
                                .difference(&transitive_cherrypicked_from)
                                .copied()
                                .copied()
                                .collect();
                            if difference.len() == 1 && difference[0] == node2 {
                                // if cherrypicked_from(node1) \ cherrypicked_from(node2) == {node2}, then
                                // we can reduce cherrypicked_from(node1) to {node2} without losing
                                // information.
                                reduced_cherrypicked_from = vec![node2];
                                break;
                            } else {
                                let new_reduced = difference;
                                if new_reduced.len() < reduced_cherrypicked_from.len() {
                                    reduced_cherrypicked_from = new_reduced;
                                }
                            }
                        }
                    }
                    if reduced_cherrypicked_from.len() == 1 {
                        stats.reducible_cherrypickedfrom += 1;
                    } else if reduced_cherrypicked_from.len() < cherrypicked_from.len() {
                        stats.partially_reducible_cherrypickedfrom += 1;
                        log::debug!(
                            "{} is not reducible to being cherry-picked from a single commit. However, it can be partially reduced.",
                            graph.properties().swhid(node1)
                        )
                    } else {
                        log::debug!(
                            "{} is not reducible to being cherry-picked from a single commit.",
                            graph.properties().swhid(node1)
                        )
                    }
                }

                reduced_cherrypicks_to_originals.push((node1, reduced_cherrypicked_from));
                (stats, reduced_cherrypicks_to_originals)
            },
        )
        .reduce(
            || (ReductionStats::default(), Vec::new()),
            |(stats1, mut cp1), (stats2, mut cp2)| {
                if cp1.len() < cp2.len() {
                    (cp1, cp2) = (cp2, cp1)
                }
                cp1.extend(cp2);
                (stats1 + stats2, cp1)
            },
        );

    (
        stats,
        reduced_cherrypicks_to_originals.into_iter().collect(),
    )
}
