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

use anyhow::Result;
use dataset_writer::StructArrayBuilder;
use dataset_writer::arrow::array::*;
use dataset_writer::arrow::datatypes::DataType::*;
use dataset_writer::arrow::datatypes::*;
use dataset_writer::parquet::basic::{Compression, Encoding, ZstdLevel};
use dataset_writer::parquet::file::properties::EnabledStatistics;
use dataset_writer::parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
use osv::schema::RangeType;
use rapidhash::RapidHashSet;
use swh_graph::graph::SwhGraph;

use crate::{Stats, git, reference_metrics};

pub fn schema() -> Schema {
    Schema::new(vec![
        Field::new("revision_id", UInt64, false),
        Field::new(
            "vulnerability_filename",
            Dictionary(Int32.into(), Utf8.into()),
            false,
        ),
    ])
}

pub fn writer_properties<G: SwhGraph>(graph: &G) -> WriterPropertiesBuilder {
    WriterProperties::builder()
        // We sort all revisions for a given same vulnerability.
        // swh-graph's LLP also makes it likely they are close numbers
        .set_column_encoding("revision_id".into(), Encoding::DELTA_BINARY_PACKED)
        .set_column_statistics_enabled("revision_id".into(), EnabledStatistics::Page)
        .set_column_bloom_filter_enabled("revision_id".into(), true)
        .set_column_compression(
            "revision_id".into(),
            Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
        )
        // Textual data, with plenty of repetition, but it's encoded as a dictionary
        // so each chunk will only have a small number of integers
        //
        // implicit and the method doesn't exist anyway:
        //.dictionary_data_page_encoding(
        //    "vulnerability_filename".into(),
        //    Encoding::RLE_DICTIONARY,
        //)
        .set_column_statistics_enabled("vulnerability_filename".into(), EnabledStatistics::Page)
        .set_column_bloom_filter_enabled("vulnerability_filename".into(), true)
        .set_column_compression(
            "vulnerability_filename".into(),
            Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
        )
        .set_key_value_metadata(Some(crate::parquet_metadata(graph)))
}

pub struct Builder {
    pub revision_id: UInt64Builder,
    pub vulnerability_filename: StringDictionaryBuilder<Int32Type>,
}

impl Default for Builder {
    fn default() -> Self {
        Builder {
            revision_id: UInt64Builder::new_from_buffer(
                Default::default(),
                None, // Values are not nullable -> validity buffer not needed
            ),
            vulnerability_filename: StringDictionaryBuilder::default(), // TODO: don't use validity buffer,
        }
    }
}

impl StructArrayBuilder for Builder {
    fn len(&self) -> usize {
        self.revision_id.len()
    }

    fn buffer_size(&self) -> usize {
        // No validity slice
        self.revision_id.len() * 8
            // TODO: count values in the dictionary too
         + self.vulnerability_filename.len() * 4 // StringDictionaryBuilder<Int32Type> -> 4-bytes indices
+ self.vulnerability_filename.validity_slice().map(|s| s.len()).unwrap_or(0)
    }

    fn finish(&mut self) -> Result<StructArray> {
        let columns: Vec<Arc<dyn Array>> = vec![
            Arc::new(self.revision_id.finish()),
            Arc::new(self.vulnerability_filename.finish()),
        ];

        Ok(StructArray::new(
            schema().fields().clone(),
            columns,
            None, // nulls
        ))
    }
}

pub fn list_commits_affected_by_vulnerability<G, CB, NMM, BorrowedNMM>(
    filename: String,
    doc: osv::schema::Vulnerability,
    git_range_parser: &git::GitRangeParser<G, CB, NMM, BorrowedNMM>,
    stats: &Stats,
    reference_metrics: &Option<HashMap<String, Vec<reference_metrics::ReferenceMetric>>>,
) -> Result<Option<RapidHashSet<usize>>>
where
    G: swh_graph::graph::SwhForwardGraph
        + swh_graph::graph::SwhBackwardGraph
        + swh_graph::graph::SwhGraphWithProperties<Maps: swh_graph::properties::Maps>,
    CB: swh_graph::views::contiguous_subgraph::ContractionBackend,
    NMM: crate::multimap::NodeMultimap<swh_graph::graph::NodeId>,
    BorrowedNMM: std::borrow::Borrow<NMM>,
{
    if let Some(affected) = &doc.affected {
        let mut error = false;
        let mut has_git_ranges = false;
        let mut vulnerable_commits = RapidHashSet::with_hasher(Default::default());
        for affected in affected {
            if let Some(ranges) = &affected.ranges {
                for range in ranges {
                    match range.range_type {
                        RangeType::Git => {
                            has_git_ranges = true;
                            match git::list_commits_in_git_range(
                                &filename,
                                &doc,
                                range,
                                git_range_parser,
                                stats,
                            ) {
                                Ok(Some(commits)) => {
                                    if let Some(reference_metrics) = reference_metrics {
                                        let v = Vec::new();
                                        log::info!(
                                            "For {}, got {} commits, expected: {:?}",
                                            doc.id,
                                            commits.len(),
                                            reference_metrics
                                                .get(&doc.id)
                                                .unwrap_or(&v)
                                                .iter()
                                                .map(|reference_metric| (
                                                    reference_metric
                                                        .metadata
                                                        .old_vulnerable_revision_number,
                                                    reference_metric
                                                        .metadata
                                                        .new_vulnerable_revision
                                                        .len()
                                                ))
                                                .collect::<Vec<_>>(),
                                        );
                                    }
                                    for commit in commits {
                                        vulnerable_commits.insert(commit);
                                    }
                                    stats
                                        .total_successful_ranges
                                        .fetch_add(1, Ordering::Relaxed);
                                }
                                Ok(None) => error = true,
                                Err(e) => log::warn!("Could not process {filename}: {e}"),
                            }
                        }
                        RangeType::Semver => log::debug!("TODO: semver range"),
                        RangeType::Ecosystem => log::debug!("TODO: ecosystem range"),
                        _ => log::warn!("Unknown range type: {:?}", range.range_type),
                    }
                }
            }
        }
        stats.total_vulnerable_commits_x_vulnerabilities.fetch_add(
            u64::try_from(vulnerable_commits.len())
                .expect("number of vulnerable commits overflowed u64"),
            Ordering::Relaxed,
        );
        if !error && has_git_ranges {
            stats
                .total_successful_documents_with_git_ranges
                .fetch_add(1, Ordering::Relaxed);
        }
        if !has_git_ranges {
            stats
                .total_documents_without_git_ranges
                .fetch_add(1, Ordering::Relaxed);
        }
        Ok(Some(vulnerable_commits))
    } else {
        Ok(None)
    }
}
