// 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::collections::HashMap;
use std::path::Path;

use anyhow::{Context, Result, bail};
use epserde::deser::{Deserialize, Flags, MemCase};
use swh_graph::graph::{NodeId, SwhGraphWithProperties};
use swh_graph::properties;
use swh_graph::{NodeType, SWHID};
use swh_graph_stdlib::connectivity::SubgraphWccs;

use crate::VulnerabilityId;
use crate::multimap::{BvNodeMultimap, NodeMultimap};

pub struct Commit2Vuln {
    id2filename: HashMap<VulnerabilityId, String>,
    commit2vuln: BvNodeMultimap,
    wccs: Option<MemCase<SubgraphWccs>>,
}

impl Commit2Vuln {
    pub fn new(sqlite_path: &Path, commit2vuln_path: &Path) -> Result<Self> {
        let filename2id = crate::db::read_filenames_from_sqlite(sqlite_path)
            .context("Could not read vulnerability filenames from sqlite")?;
        let id2filename = filename2id
            .into_iter()
            .map(|(filename, id)| (id, filename))
            .collect();

        let commit2vuln =
            BvNodeMultimap::mmap(commit2vuln_path).context("Could not mmap commit2vuln")?;

        Ok(Commit2Vuln {
            id2filename,
            commit2vuln,
            wccs: None,
        })
    }

    pub fn with_subgraphwccs(&mut self, path: &Path) -> Result<&mut Self> {
        self.wccs = Some(
            unsafe { <SubgraphWccs>::mmap(path, Flags::RANDOM_ACCESS) }
                .with_context(|| format!("Could not mmap SubgraphWccs from {}", path.display()))?,
        );
        Ok(self)
    }

    pub fn get_commit_vuln_ids(
        &self,
        mut commit: NodeId,
    ) -> Result<Option<impl Iterator<Item = VulnerabilityId> + '_>> {
        if let Some(wccs) = &self.wccs {
            match wccs.uncase().contraction().node_id_from_underlying(commit) {
                Some(compacted_commit) => commit = compacted_commit,
                None => return Ok(None),
            }
        }

        Ok(Some(self.commit2vuln.get(commit).into_iter().map(
            |vuln_id| {
                VulnerabilityId(u64::try_from(vuln_id).expect("Vulnerability id overflowed u64"))
            },
        )))
    }

    pub fn get_commit_vuln_filenames(
        &self,
        commit: NodeId,
    ) -> Result<Option<impl Iterator<Item = (VulnerabilityId, &str)>>> {
        match self.get_commit_vuln_ids(commit) {
            Ok(Some(it)) => Ok(Some(it.map(|vuln_id| {
                (
                    vuln_id,
                    self.id2filename
                        .get(&vuln_id)
                        .expect("Vulnerability id in commit2vuln does not exist in id2filename")
                        .as_ref(),
                )
            }))),
            Ok(None) => Ok(None),
            Err(e) => Err(e),
        }
    }
}

#[derive(serde::Serialize)]
struct OutputRecord<'a> {
    swhid: SWHID,
    vuln_id: VulnerabilityId,
    vuln_filename: &'a str,
}

pub fn write_vulnerabilities_for_swhids<G, I, W>(
    commit2vuln: &mut Commit2Vuln,
    graph: &G,
    swhids: I,
    writer: &mut csv::Writer<W>,
) -> Result<()>
where
    G: SwhGraphWithProperties<Maps: properties::Maps>,
    I: Iterator<Item = Result<SWHID>>,
    W: std::io::Write,
{
    for (i, swhid_result) in swhids.enumerate() {
        let swhid = swhid_result.with_context(|| format!("Invalid SWHID on line {i}"))?;

        let Ok(node) = graph.properties().node_id(swhid) else {
            log::warn!("Unknown SWHID: {swhid}");
            continue;
        };

        if !matches!(
            graph.properties().node_type(node),
            NodeType::Revision | NodeType::Release
        ) {
            bail!("{swhid} is neither a revision nor a release");
        }

        if let Some(vulns) = commit2vuln
            .get_commit_vuln_filenames(node)
            .with_context(|| format!("Could not get vulnerabilities of {swhid}"))?
        {
            for (vuln_id, vuln_filename) in vulns {
                writer
                    .serialize(OutputRecord {
                        swhid,
                        vuln_id,
                        vuln_filename,
                    })
                    .context("Could not write to stdout")?;
            }
            writer.flush().context("Could not flush stdout")?;
        }
    }

    Ok(())
}
