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

use anyhow::{Context, 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 dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use rayon::prelude::*;

use crate::VulnerabilityId;
use crate::db;

pub fn sqlite_to_parquet(sqlite_path: &Path, parquet_dir: &Path) -> Result<()> {
    let vulnerability_ids: HashMap<_, _> = db::read_filenames_from_sqlite(sqlite_path)?
        .into_iter()
        .map(|(filename, id)| (id, filename))
        .collect();

    log::info!("Sorting vulnerabilities...");
    let mut vulnerability_ids: Vec<_> = vulnerability_ids.into_iter().collect();
    vulnerability_ids.par_sort_unstable();
    let num_vulnerabilities = vulnerability_ids.len();

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(num_vulnerabilities),
    );
    pl.start(format!(
        "Writing {num_vulnerabilities} vulnerabilities to Parquet..."
    ));

    let dataset_writer = dataset_writer::ParallelDatasetWriter::<
        dataset_writer::ParquetTableWriter<Builder>,
    >::with_schema(
        parquet_dir.to_owned(),
        (Arc::new(schema()), writer_properties().build()),
    )
    .with_context(|| format!("Could not create {}", parquet_dir.display()))?;

    vulnerability_ids.into_par_iter().try_for_each_init(
        || {
            (
                pl.clone(),
                dataset_writer
                    .get_thread_writer()
                    .expect("Could not get thread writer"),
            )
        },
        |(pl, writer), (VulnerabilityId(vuln_id), vuln_filename)| {
            let builder = writer.builder()?;
            builder.ids.append_value(vuln_id);
            builder.filenames.append_value(vuln_filename);
            pl.light_update();
            Ok(())
        },
    )
}
pub struct Builder {
    pub ids: UInt64Builder,
    pub filenames: StringBuilder,
}

impl Default for Builder {
    fn default() -> Self {
        Builder {
            ids: UInt64Builder::new_from_buffer(
                Default::default(),
                None, // not nullable -> validity buffer not needed
            ),
            filenames: StringBuilder::new(), // FIXME: exclude validity buffer, values can't be null
        }
    }
}

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

    fn buffer_size(&self) -> usize {
        // No validity slice
        self.ids.len() * 8
            + self.filenames.offsets_slice().len() * 4
            + self.filenames.values_slice().len()
            + self
                .filenames
                .validity_slice()
                .map(|slice| slice.len())
                .unwrap_or(0)
    }

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

        Ok(StructArray::new(
            schema().fields().clone(),
            columns,
            None, // nulls
        ))
    }
}
pub fn schema() -> Schema {
    Schema::new(vec![
        Field::new("vulnerability_id", UInt64, false),
        Field::new("vulnerability_filename", Utf8, false),
    ])
}

pub fn writer_properties() -> WriterPropertiesBuilder {
    WriterProperties::builder()
        // We sort vulnerabilities
        .set_column_encoding("vulnerability_id".into(), Encoding::DELTA_BINARY_PACKED)
        .set_column_statistics_enabled("vulnerability_id".into(), EnabledStatistics::Page)
        .set_column_bloom_filter_enabled("vulnerability_id".into(), true)
        .set_column_compression(
            "vulnerability_id".into(),
            Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
        )
        // Textual data, somewhat sorted
        .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()),
        )
}
