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

#![cfg(feature = "arrow")]

use std::collections::HashMap;
use std::num::NonZeroU16;
use std::sync::Arc;

use anyhow::Result;
use dataset_writer::StructArrayBuilder;
use dataset_writer::arrow::array::RecordBatch;
use dataset_writer::parquet::arrow::ArrowWriter;
use dsi_bitstream::prelude::BE;
use rusqlite::Connection;
use swh_graph::graph::SwhGraph;
use swh_graph::swhid;
use webgraph_cli::build::ef::{CliArgs, build_elias_fano};

use swh_osv::commit2vuln::Commit2Vuln;
use swh_osv::commit2vuln_parquet_to_bvgraph::build_commit_to_vuln_bvgraph;
use swh_osv::multimap::{BvNodeMultimap, NodeMultimap};
use swh_osv::{VulnerabilityId, vuln2commit_parquet};

#[path = "./graphs.rs"]
mod graphs;

const NO_SUBGRAPH: Option<&swh_graph_stdlib::connectivity::SubgraphWccs<Vec<usize>, Vec<usize>>> =
    None;

fn build_ef_index(basename: &std::path::Path) -> Result<()> {
    build_elias_fano::<BE>(
        webgraph_cli::GlobalArgs { log_interval: None },
        CliArgs {
            basename: basename.to_path_buf(),
            number_of_nodes: None,
        },
    )
}

struct ParquetTestFixture {
    _temp_dir: tempfile::TempDir,
    parquet_dir: std::path::PathBuf,
    output_path: std::path::PathBuf,
}

impl ParquetTestFixture {
    fn new() -> Result<Self> {
        let temp_dir = tempfile::tempdir()?;
        let parquet_dir = temp_dir.path().join("parquet");
        std::fs::create_dir(&parquet_dir)?;
        let output_path = temp_dir.path().join("commit2vuln");

        Ok(Self {
            _temp_dir: temp_dir,
            parquet_dir,
            output_path,
        })
    }

    fn add_parquet_file(&self, filename: &str, data: Vec<(u64, String)>) -> Result<()> {
        let path = self.parquet_dir.join(filename);
        let schema = Arc::new(vuln2commit_parquet::schema());

        let file = std::fs::File::create(path)?;
        let mut writer = ArrowWriter::try_new(file, schema.clone(), None)?;

        let mut builder = vuln2commit_parquet::Builder::default();

        for (revision_id, vulnerability_filename) in data {
            builder.revision_id.append_value(revision_id);
            builder
                .vulnerability_filename
                .append_value(&vulnerability_filename);
        }

        let struct_array = builder.finish()?;
        let batch = RecordBatch::from(struct_array);

        writer.write(&batch)?;
        writer.close()?;

        Ok(())
    }
}

#[test]
fn test_build_commit_to_vuln_bvgraph_simple() -> Result<()> {
    let fixture = ParquetTestFixture::new()?;

    fixture.add_parquet_file("data1.parquet", vec![(1, "vuln1.json".to_string())])?;
    fixture.add_parquet_file(
        "data2.parquet",
        vec![(2, "vuln1.json".to_string()), (3, "vuln2.json".to_string())],
    )?;

    let mut vulnerability_ids = HashMap::new();
    vulnerability_ids.insert("vuln1.json".to_string(), VulnerabilityId(0));
    vulnerability_ids.insert("vuln2.json".to_string(), VulnerabilityId(1));

    let graph = graphs::get_10rev_linear_graph();
    let num_nodes = graph.num_nodes();

    build_commit_to_vuln_bvgraph(
        num_nodes,
        &vulnerability_ids,
        &graph,
        NO_SUBGRAPH,
        &fixture.parquet_dir,
        NonZeroU16::new(1).unwrap(),
        &fixture.output_path,
    )?;

    build_ef_index(&fixture.output_path)?;

    let multimap = BvNodeMultimap::mmap(&fixture.output_path)?;

    let vuln_ids: Vec<_> = multimap.get(1).into_iter().collect();
    assert_eq!(vuln_ids, vec![0]);

    let vuln_ids: Vec<_> = multimap.get(2).into_iter().collect();
    assert_eq!(vuln_ids, vec![0]);

    let vuln_ids: Vec<_> = multimap.get(3).into_iter().collect();
    assert_eq!(vuln_ids, vec![1]);

    let vuln_ids: Vec<_> = multimap.get(4).into_iter().collect();
    assert_eq!(vuln_ids, Vec::<usize>::new());

    Ok(())
}

#[test]
fn test_build_commit_to_vuln_bvgraph_multiple_vulns_per_commit() -> Result<()> {
    let fixture = ParquetTestFixture::new()?;

    fixture.add_parquet_file(
        "data.parquet",
        vec![
            (1, "vuln1.json".to_string()),
            (1, "vuln2.json".to_string()),
            (1, "vuln3.json".to_string()),
        ],
    )?;

    let mut vulnerability_ids = HashMap::new();
    vulnerability_ids.insert("vuln1.json".to_string(), VulnerabilityId(0));
    vulnerability_ids.insert("vuln2.json".to_string(), VulnerabilityId(1));
    vulnerability_ids.insert("vuln3.json".to_string(), VulnerabilityId(2));

    let graph = graphs::get_10rev_linear_graph();
    let num_nodes = graph.num_nodes();

    build_commit_to_vuln_bvgraph(
        num_nodes,
        &vulnerability_ids,
        &graph,
        NO_SUBGRAPH,
        &fixture.parquet_dir,
        NonZeroU16::new(1).unwrap(),
        &fixture.output_path,
    )?;

    build_ef_index(&fixture.output_path)?;

    let multimap = BvNodeMultimap::mmap(&fixture.output_path)?;

    let vuln_ids: Vec<_> = multimap.get(1).into_iter().collect();
    assert_eq!(vuln_ids, vec![0, 1, 2]);

    Ok(())
}

#[test]
fn test_build_commit_to_vuln_bvgraph_empty_parquet() -> Result<()> {
    let fixture = ParquetTestFixture::new()?;

    fixture.add_parquet_file("data.parquet", vec![])?;

    let vulnerability_ids = HashMap::new();
    let graph = graphs::get_10rev_linear_graph();
    let num_nodes = graph.num_nodes();

    build_commit_to_vuln_bvgraph(
        num_nodes,
        &vulnerability_ids,
        &graph,
        NO_SUBGRAPH,
        &fixture.parquet_dir,
        NonZeroU16::new(1).unwrap(),
        &fixture.output_path,
    )?;

    build_ef_index(&fixture.output_path)?;

    let multimap = BvNodeMultimap::mmap(&fixture.output_path)?;

    let vuln_ids: Vec<_> = multimap.get(1).into_iter().collect();
    assert_eq!(vuln_ids, Vec::<usize>::new());

    Ok(())
}

fn create_test_db_with_filenames(filenames: Vec<&str>) -> Result<tempfile::NamedTempFile> {
    let temp_file = tempfile::NamedTempFile::new()?;
    let conn = Connection::open(temp_file.path())?;

    conn.execute(
        "CREATE TABLE vulnerabilities (
            id INTEGER PRIMARY KEY,
            filename TEXT NOT NULL,
            doc TEXT NOT NULL
        )",
        [],
    )?;

    for (i, filename) in filenames.into_iter().enumerate() {
        conn.execute(
            "INSERT INTO vulnerabilities (id, filename, doc) VALUES (?, ?, ?)",
            [&(i as u64).to_string(), filename, "{}"],
        )?;
    }

    Ok(temp_file)
}

#[test]
fn test_write_vulnerabilities_for_swhids() -> Result<()> {
    let fixture = ParquetTestFixture::new()?;

    fixture.add_parquet_file(
        "data.parquet",
        vec![
            (1, "vuln1.json".to_string()),
            (2, "vuln1.json".to_string()),
            (2, "vuln2.json".to_string()),
        ],
    )?;

    let mut vulnerability_ids = HashMap::new();
    vulnerability_ids.insert("vuln1.json".to_string(), VulnerabilityId(0));
    vulnerability_ids.insert("vuln2.json".to_string(), VulnerabilityId(1));

    let graph = graphs::get_10rev_linear_graph();
    let num_nodes = graph.num_nodes();

    build_commit_to_vuln_bvgraph(
        num_nodes,
        &vulnerability_ids,
        &graph,
        NO_SUBGRAPH,
        &fixture.parquet_dir,
        NonZeroU16::new(1).unwrap(),
        &fixture.output_path,
    )?;

    build_ef_index(&fixture.output_path)?;

    let db_file = create_test_db_with_filenames(vec!["vuln1.json", "vuln2.json"])?;
    let mut commit2vuln = Commit2Vuln::new(db_file.path(), &fixture.output_path)?;

    let swhids = vec![
        Ok(swhid!(swh:1:rev:0000000000000000000000000000000000000001)),
        Ok(swhid!(swh:1:rev:0000000000000000000000000000000000000099)), // unknown, should be skipped
        Ok(swhid!(swh:1:rev:0000000000000000000000000000000000000002)),
    ];

    let mut output = Vec::new();
    let mut writer = csv::Writer::from_writer(&mut output);

    swh_osv::commit2vuln::write_vulnerabilities_for_swhids(
        &mut commit2vuln,
        &graph,
        swhids.into_iter(),
        &mut writer,
    )?;

    drop(writer);

    assert_eq!(
        String::from_utf8(output)?,
        "swhid,vuln_id,vuln_filename\n\
         swh:1:rev:0000000000000000000000000000000000000001,0,vuln1.json\n\
         swh:1:rev:0000000000000000000000000000000000000002,0,vuln1.json\n\
         swh:1:rev:0000000000000000000000000000000000000002,1,vuln2.json\n"
    );

    Ok(())
}
