// 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::marker::PhantomData;
use std::sync::atomic::Ordering;

use anyhow::{Context, Result, bail, ensure};
use osv::schema::{Event, Range, Vulnerability};
use rapidhash::RapidHashSet;
use sux::dict::elias_fano::EfSeqDict;
use swh_graph::graph::{
    NodeId, SwhBackwardGraph, SwhForwardGraph, SwhGraphWithProperties, SwhLabeledBackwardGraph,
};
use swh_graph::labels::Branch;
use swh_graph::properties;
use swh_graph::properties::NodeIdFromSwhidError;
use swh_graph::views::contiguous_subgraph::{Contraction, ContractionBackend};
use swh_graph::views::{Subgraph, Transposed};
use swh_graph::{NodeType, SWHID};
use swh_graph_stdlib::collections::{NodeSet, ReadNodeSet, SmallNodeSet};
use swh_graph_stdlib::iter_nodes;

use crate::multimap::*;
use crate::traversals;
use crate::{EventType, Stats};

pub struct GitRangeParser<
    G: SwhGraphWithProperties<Maps: properties::Maps>,
    CB: ContractionBackend = EfSeqDict,
    NMM: NodeMultimap<NodeId> = EmptyNodeMultimap,
    BorrowedNNM: Borrow<NMM> = NMM,
> {
    graph: G,
    nodeid_contraction: Option<Contraction<CB>>,
    originals_to_cherrypicks: Option<BorrowedNNM>,
    cherrypicks_to_originals: Option<BorrowedNNM>,
    marker: PhantomData<NMM>,
}

impl<G: SwhGraphWithProperties<Maps: properties::Maps>> GitRangeParser<G> {
    #[allow(clippy::type_complexity)] // we don't have a choice
    pub fn new(
        graph: G,
    ) -> GitRangeParser<Subgraph<G, impl Fn(NodeId) -> bool, fn(usize, usize) -> bool>>
    where
        G: Clone,
    {
        GitRangeParser {
            graph: Subgraph::with_node_constraint(graph, "rev".parse().unwrap()),
            nodeid_contraction: None,
            originals_to_cherrypicks: None,
            cherrypicks_to_originals: None,
            marker: PhantomData,
        }
    }
}

impl<
    G: SwhGraphWithProperties<Maps: properties::Maps>,
    CB: ContractionBackend,
    NMM: NodeMultimap<NodeId>,
    BorrowedNNM: Borrow<NMM>,
> GitRangeParser<G, CB, NMM, BorrowedNNM>
{
    pub fn graph(&self) -> &G {
        &self.graph
    }

    pub fn with_nodeid_contraction<CB2: ContractionBackend>(
        self,
        nodeid_contraction: Option<Contraction<CB2>>,
    ) -> Result<GitRangeParser<G, CB2, NMM, BorrowedNNM>> {
        let GitRangeParser {
            graph,
            nodeid_contraction: _,
            originals_to_cherrypicks,
            cherrypicks_to_originals,
            marker: PhantomData,
        } = self;
        Ok(GitRangeParser {
            graph,
            nodeid_contraction,
            originals_to_cherrypicks,
            cherrypicks_to_originals,
            marker: PhantomData,
        })
    }

    pub fn with_node_multimaps<NMM2: NodeMultimap<NodeId>, BorrowedNNM2: Borrow<NMM2>>(
        self,
        originals_to_cherrypicks: Option<BorrowedNNM2>,
        cherrypicks_to_originals: Option<BorrowedNNM2>,
    ) -> Result<GitRangeParser<G, CB, NMM2, BorrowedNNM2>> {
        let GitRangeParser {
            graph,
            nodeid_contraction,
            originals_to_cherrypicks: _,
            cherrypicks_to_originals: _,
            marker: PhantomData,
        } = self;
        Ok(GitRangeParser {
            graph,
            nodeid_contraction,
            originals_to_cherrypicks,
            cherrypicks_to_originals,
            marker: PhantomData,
        })
    }

    /// Returns a an array mapping EventType to a set of events of each type, and a boolean indicating
    /// if there is a "0" pseudo-commit in the "introduced" set; or `None` in case of expected error
    /// (counted in stats)
    pub fn parse_git_range(
        &self,
        range: &Range,
        stats: Option<&Stats>,
    ) -> Result<Option<([SmallNodeSet; 4], bool)>> {
        let mut commits_by_event_type = [
            SmallNodeSet::default(),
            SmallNodeSet::default(),
            SmallNodeSet::default(),
            SmallNodeSet::default(),
        ];
        let mut has_zero_commit = false;
        for event in &range.events {
            let (event_type, commit_id) = match event {
                Event::Introduced(commit_id) => (EventType::Introduced, commit_id),
                Event::Fixed(commit_id) => (EventType::Fixed, commit_id),
                Event::LastAffected(commit_id) => (EventType::LastAffected, commit_id),
                Event::Limit(commit_id) => (EventType::Limit, commit_id),
                _ => bail!("Unexpected event type: {:?}", range.events),
            };
            if commit_id == "0" {
                ensure!(
                    event_type == EventType::Introduced,
                    "Commit id \"0\" is only valid for \"introduced\" events, but it is used in {event_type:?}"
                );
                has_zero_commit = true;
                continue;
            }
            let mut hash = [0u8; 20];
            if let Err(e) = faster_hex::hex_decode(commit_id.as_bytes(), &mut hash) {
                log::debug!("Invalid commit id ({commit_id}): {e}");
                if let Some(stats) = stats {
                    stats
                        .total_invalid_commit_ids
                        .fetch_add(1, Ordering::Relaxed);
                }
                return Ok(None);
            }
            let node_id = match self.graph.properties().node_id(SWHID {
                namespace_version: 1,
                node_type: NodeType::Revision,
                hash,
            }) {
                Ok(node_id) => node_id,
                Err(NodeIdFromSwhidError::UnknownSwhid(swhid)) => {
                    log::debug!("Unknown SWHID: {swhid}");
                    if let Some(stats) = stats {
                        stats.total_unknown_commits.fetch_add(1, Ordering::Relaxed);
                    }
                    return Ok(None);
                }
                Err(NodeIdFromSwhidError::InternalError(e)) => {
                    bail!("Could not get node id for commit {commit_id}: {e}")
                }
            };

            // Find all nodes that 'node_id' is a cherry-pick of, and all nodes that it cherry-picks;
            // transitively.
            let mut node_ids = RapidHashSet::with_hasher(Default::default());

            // translate node ids to the compacted ids used in cherrypicks_to_originals and
            // originals_to_cherrypicks
            let mut node_ids_to_process = if let Some(nodeid_contraction) = &self.nodeid_contraction
            {
                vec![nodeid_contraction.node_id_from_underlying(node_id).with_context(|| format!("Commit {commit_id} has node id {node_id} which is missing from the nodeid contraction"))?]
            } else {
                vec![node_id]
            };
            node_ids.extend(node_ids_to_process.iter().copied());
            while let Some(node_id) = node_ids_to_process.pop() {
                if let Some(cherrypicks_to_originals) = &self.cherrypicks_to_originals {
                    for original in cherrypicks_to_originals.borrow().get(node_id) {
                        let new = node_ids.insert(original);
                        if !new {
                            node_ids_to_process.push(original);
                        }
                    }
                }
                if let Some(originals_to_cherrypicks) = &self.originals_to_cherrypicks {
                    for cherrypick in originals_to_cherrypicks.borrow().get(node_id) {
                        let new = node_ids.insert(cherrypick);
                        if !new {
                            node_ids_to_process.push(cherrypick);
                        }
                    }
                }
            }
            // translate node ids back from the compacted numbers to the originals
            if let Some(nodeid_contraction) = &self.nodeid_contraction {
                node_ids = node_ids
                    .into_iter()
                    .map(|compact_node_id| nodeid_contraction.underlying_node_id(compact_node_id))
                    .collect()
            }

            for node_id in node_ids {
                commits_by_event_type[event_type as usize].insert(node_id);
            }
        }

        Ok(Some((commits_by_event_type, has_zero_commit)))
    }
}

/// Adds the given range to statistics, and returns vulnerable commits
///
/// Returns `Ok(Some(_))` if all went well, `Ok(None)` in case of expected error
/// (counted in stats), or `Err(_)` in case of unexpected error (to be logged by caller)
pub fn list_commits_in_git_range<
    G: SwhForwardGraph + SwhBackwardGraph + SwhGraphWithProperties<Maps: properties::Maps>,
    CB: ContractionBackend,
    NMM: NodeMultimap<NodeId>,
    BorrowedNNM: Borrow<NMM>,
>(
    filename: &str,
    doc: &Vulnerability,
    range: &Range,
    git_range_parser: &GitRangeParser<G, CB, NMM, BorrowedNNM>,
    stats: &Stats,
) -> Result<Option<RapidHashSet<NodeId>>> {
    let (mut commits_by_event_type, has_zero_commit) = match git_range_parser
        .parse_git_range(range, Some(stats))
        .with_context(|| format!("Could not parse git range in {filename}"))?
    {
        Some(res) => res,
        None => return Ok(None),
    };
    if has_zero_commit {
        // add every root revision of revisions mentioned in other event types to the "introduced"
        // set
        let Ok(()) = traversals::for_each_root(
            &git_range_parser.graph(),
            commits_by_event_type
                .iter()
                .flat_map(|set| set.iter())
                .collect(),
            |root| -> Result<_, std::convert::Infallible> {
                commits_by_event_type[EventType::Introduced as usize].insert(root);
                Ok(())
            },
        );
    }

    let mut ignore_nodes = RapidHashSet::with_hasher(Default::default());
    for event_type in [EventType::Fixed, EventType::LastAffected, EventType::Limit] {
        ignore_nodes.extend(&commits_by_event_type[event_type as usize]);
    }
    // List all commits that are successors of the introduction.
    // Ignore commits that are only reachable through a "fixed", "last_affected", or
    // "limit" commit as an optimization (because we easily know they are already patched)
    let coreachable_by_introduction_not_through_fix: RapidHashSet<_> = iter_nodes(
        &Transposed(Subgraph::with_node_filter(
            git_range_parser.graph(),
            |node| !ignore_nodes.contains(&node),
        )),
        &commits_by_event_type[EventType::Introduced as usize],
    )
    .collect();
    drop(ignore_nodes); // free memory

    // List commits that are patched, among the candidates we have so far.
    // Skip commits that are not reachable from an "introduced" commit as an optimization
    // (because we already know they are not vulnerable, so there is nothing to patch)
    let mut fix_commits = commits_by_event_type[EventType::Fixed as usize].clone();
    for last_affected in &commits_by_event_type[EventType::LastAffected as usize] {
        fix_commits.insert(last_affected);
    }
    let coreachable_by_fix: RapidHashSet<_> = iter_nodes(
        &Transposed(Subgraph::with_node_filter(
            git_range_parser.graph(),
            |node| coreachable_by_introduction_not_through_fix.contains(&node),
        )),
        &fix_commits,
    )
    .collect();
    let mut vulnerable = if commits_by_event_type[EventType::Limit as usize].is_empty() {
        // vulnerable = coreachable_by_introduction_not_through_fix - coreachable_by_fix
        let mut vulnerable = RapidHashSet::with_capacity_and_hasher(
            coreachable_by_introduction_not_through_fix
                .len()
                .saturating_sub(coreachable_by_fix.len()),
            Default::default(),
        );
        for node in &coreachable_by_introduction_not_through_fix {
            if !coreachable_by_fix.contains(node) {
                vulnerable.insert(*node);
            }
        }
        vulnerable
    } else {
        let reachable_by_limit: RapidHashSet<_> = iter_nodes(
            git_range_parser.graph(),
            &commits_by_event_type[EventType::Limit as usize],
        )
        .collect();

        // vulnerable = (
        //      (coreachable_by_introduction_not_through_fix - coreachable_by_fix)
        //      /\ reachable_by_limit)
        // ) - limit
        //
        // this implies that every commit in 'vulnerable' is on a path from an introduction commit
        // to a limit commit.
        let mut vulnerable =
            RapidHashSet::with_capacity_and_hasher(reachable_by_limit.len(), Default::default());
        for node in &coreachable_by_introduction_not_through_fix {
            if !coreachable_by_fix.contains(node)
                && reachable_by_limit.contains(node)
                && !commits_by_event_type[EventType::Limit as usize].contains(*node)
            {
                vulnerable.insert(*node);
            }
        }
        vulnerable
    };

    // add last_affected to the set of vulnerable commits, as they are vulnerable by definition
    for node in &commits_by_event_type[EventType::LastAffected as usize] {
        vulnerable.insert(node);
    }

    stats
        .total_vulnerable_commits_x_vulnerabilities_x_ranges
        .fetch_add(
            u64::try_from(vulnerable.len()).expect("number of vulnerable commits overflowed u64"),
            Ordering::Relaxed,
        );

    log::debug!(
        "For {}: {:?} introduced, {} fixed, {} last_affected, {} limit. has_zero_commit={has_zero_commit}. {} coreachable_by_introduction_not_through_fix, {} fix_commits, {} coreachable_by_fix, {} vulnerable",
        doc.id,
        commits_by_event_type[EventType::Introduced as usize].len(),
        commits_by_event_type[EventType::Fixed as usize].len(),
        commits_by_event_type[EventType::LastAffected as usize].len(),
        commits_by_event_type[EventType::Limit as usize].len(),
        coreachable_by_introduction_not_through_fix.len(),
        fix_commits.len(),
        coreachable_by_fix.len(),
        vulnerable.len()
    );
    Ok(Some(vulnerable))
}

/// Extracts tag names that point to a given revision.
///
/// Either through Release nodes or through Snapshot branches with "refs/tags/" prefix.
pub fn get_tag_names<
    G: SwhLabeledBackwardGraph
        + SwhGraphWithProperties<
            Maps: properties::Maps,
            Strings: properties::Strings,
            LabelNames: properties::LabelNames,
        >,
>(
    graph: &G,
    rev: NodeId,
) -> Vec<String> {
    graph
        .untyped_labeled_predecessors(rev)
        .into_iter()
        .flat_map(|(pred, labels)| -> Vec<String> {
            match graph.properties().node_type(pred) {
                NodeType::Release => match graph.properties().tag_name(pred) {
                    Some(tag_name) => vec![String::from_utf8_lossy(&tag_name).into_owned()],
                    None => vec![],
                },
                NodeType::Snapshot => labels
                    .into_iter()
                    .flat_map(|label| {
                        let label = Branch::from(label);
                        let branch_name = graph.properties().label_name(label.label_name_id());
                        let branch_name = String::from_utf8_lossy(&branch_name);
                        branch_name
                            .strip_prefix("refs/tags/") // returns None if not prefixed
                            .map(|tag_name| tag_name.to_owned())
                    })
                    .collect(),
                _ => vec![],
            }
        })
        .collect::<RapidHashSet<_>>() // dedup
        .into_iter()
        .collect()
}

/// Extracts all tag names for multiple revisions.
pub fn get_all_tag_names<
    G: SwhLabeledBackwardGraph
        + SwhGraphWithProperties<
            Maps: properties::Maps,
            Strings: properties::Strings,
            LabelNames: properties::LabelNames,
        >,
>(
    graph: &G,
    revs: impl IntoIterator<Item = NodeId>,
) -> Vec<String> {
    revs.into_iter()
        .flat_map(|rev| get_tag_names(graph, rev))
        .collect()
}

/// Filters out merge commits (revisions with multiple parents).
///
/// Returns only revisions that have at most one parent revision.
pub fn filter_non_merge_revisions<
    G: SwhForwardGraph + SwhGraphWithProperties<Maps: properties::Maps> + Sync + Send,
    I: IntoIterator<Item = NodeId>,
>(
    graph: &G,
    revs: I,
) -> Vec<NodeId> {
    use rayon::prelude::*;

    revs.into_iter()
        .collect::<Vec<_>>()
        .into_par_iter()
        .filter(|&rev| {
            let mut found_parent = false;
            for succ in graph.successors(rev) {
                match graph.properties().node_type(succ) {
                    NodeType::Directory => (),
                    NodeType::Revision => {
                        if found_parent {
                            // has two parents
                            return false;
                        } else {
                            found_parent = true;
                        }
                    }
                    _ => panic!(
                        "Revision {} has non-dir and non-rev successor: {}",
                        graph.properties().swhid(rev),
                        graph.properties().swhid(succ)
                    ),
                }
            }

            true
        })
        .collect()
}
