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

use anyhow::{Context, Result, ensure};
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 dataset_writer::{ParallelDatasetWriter, ParquetTableWriter, StructArrayBuilder};
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use lender::Lender;
use swh_graph::graph::SwhGraphWithProperties;
use swh_graph::properties;
use swh_graph::views::contiguous_subgraph::{Contraction, ContractionBackend};
use webgraph::prelude::*;

use crate::utils::IdentityContractionBackend;

pub struct BvgraphToParquet<
    G: SwhGraphWithProperties<Maps: properties::Maps> + Sync,
    N: ContractionBackend = IdentityContractionBackend,
> {
    graph: G,
    path: PathBuf,
    properties_builder: WriterPropertiesBuilder,
    fields: Vec<Field>,
    num_files: usize,
    key_column: &'static str,
    value_column: &'static str,
    key_swhids: bool,
    value_swhids: bool,
    contraction: Contraction<N>,
    uncontract_keys: bool,
    uncontract_values: bool,
}

impl<G: SwhGraphWithProperties<Maps: properties::Maps> + Sync>
    BvgraphToParquet<G, IdentityContractionBackend>
{
    pub fn new(
        graph: G,
        path: PathBuf,
        key_column: &'static str,
        value_column: &'static str,
    ) -> Self {
        let properties_builder = WriterProperties::builder()
            // keys are perfectly sorted and are often repeated
            .set_column_encoding(key_column.into(), Encoding::DELTA_BINARY_PACKED)
            .set_column_statistics_enabled(key_column.into(), EnabledStatistics::Page)
            .set_column_bloom_filter_enabled(key_column.into(), true)
            .set_column_bloom_filter_ndv(key_column.into(), 100_000) // keys are often repeated
            .set_column_compression(
                key_column.into(),
                Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
            )
            // values, sorted for a given key
            .set_column_encoding(value_column.into(), Encoding::DELTA_BINARY_PACKED)
            .set_column_statistics_enabled(value_column.into(), EnabledStatistics::Page)
            .set_column_compression(
                value_column.into(),
                Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
            )
            // other
            .set_key_value_metadata(Some(crate::parquet_metadata(&graph)));

        let fields = vec![
            Field::new(key_column, UInt64, false),
            Field::new(value_column, UInt64, false),
        ];
        Self {
            properties_builder,
            fields,
            path,
            num_files: num_cpus::get(),
            key_column,
            value_column,
            key_swhids: false,
            value_swhids: false,
            contraction: Contraction(IdentityContractionBackend {
                num_nodes: graph.num_nodes(),
            }),
            uncontract_keys: false,
            uncontract_values: false,
            graph,
        }
    }
}

impl<G: SwhGraphWithProperties<Maps: properties::Maps> + Sync, N: ContractionBackend + Sync>
    BvgraphToParquet<G, N>
{
    pub fn num_files(self, num_files: usize) -> Self {
        Self { num_files, ..self }
    }

    /// Add a column with SWHIDs of nodes in the 'key' column
    pub fn with_key_swhids(mut self) -> Self {
        if self.key_swhids {
            return self;
        }
        self.key_swhids = true;

        let name = format!("{}_swhid", self.key_column);
        self.with_swhid_column(&name)
    }

    /// Add a column with SWHIDs of nodes in the 'value' column
    pub fn with_value_swhids(mut self) -> Self {
        if self.value_swhids {
            return self;
        }
        self.value_swhids = true;

        let name = format!("{}_swhid", self.value_column);
        self.with_swhid_column(&name)
    }

    pub fn with_swhid_column(mut self, name: &str) -> Self {
        self.properties_builder = self
            .properties_builder
            .set_column_encoding(name.into(), Encoding::PLAIN)
            .set_column_statistics_enabled(name.into(), EnabledStatistics::None)
            .set_column_bloom_filter_enabled(name.into(), true)
            .set_column_bloom_filter_ndv(name.into(), 100_000) // keys are often repeated
            .set_column_compression(
                name.into(),
                Compression::ZSTD(ZstdLevel::try_new(3).unwrap()),
            );

        self.fields.push(Field::new(name, Utf8, false));

        self
    }

    /// Convert from internal node ids (eg. from connected_components.wccs) to the graph's real
    /// node ids
    pub fn with_contraction<N2: ContractionBackend>(
        self,
        contraction: Contraction<N2>,
        uncontract_keys: bool,
        uncontract_values: bool,
    ) -> Result<BvgraphToParquet<G, N2>> {
        ensure!(!self.uncontract_keys, "uncontract_keys is already set");
        ensure!(!self.uncontract_values, "uncontract_values is already set");
        let last_underlying_node_id = contraction.underlying_node_id(contraction.num_nodes() - 1);
        let expected_num_underlying_nodes = self.graph.num_nodes();
        ensure!(
            last_underlying_node_id < expected_num_underlying_nodes,
            "Contraction's underlying graph has at least {last_underlying_node_id} nodes, but given graph has {expected_num_underlying_nodes}"
        );

        let Self {
            graph,
            properties_builder,
            fields,
            path,
            num_files,
            key_column,
            value_column,
            key_swhids,
            value_swhids,
            contraction: _,
            uncontract_keys: _,
            uncontract_values: _,
        } = self;
        Ok(BvgraphToParquet {
            graph,
            properties_builder,
            fields,
            path,
            num_files,
            key_column,
            value_column,
            key_swhids,
            value_swhids,
            contraction,
            uncontract_keys,
            uncontract_values,
        })
    }

    pub fn convert(self, bvgraph: impl SplitLabeling<Label = usize>) -> Result<()> {
        let Self {
            graph,
            properties_builder,
            fields,
            path,
            num_files,
            key_column: _,
            value_column: _,
            key_swhids,
            value_swhids,
            contraction,
            uncontract_keys,
            uncontract_values,
        } = self;
        let graph = &graph;
        let schema = Arc::new(Schema::new(fields.clone()));
        let properties = properties_builder.build();
        let writer = ParallelDatasetWriter::<ParquetTableWriter<Builder>>::with_schema(
            path,
            (schema, properties),
        )
        .context("Could not create dataset")?;
        let contraction = &contraction;

        let mut pl = concurrent_progress_logger!(
            display_memory = true,
            item_name = "pair",
            local_speed = true,
            expected_updates = bvgraph.num_arcs_hint().and_then(|n| n.try_into().ok()),
        );
        pl.start("Reading .graph and writing Parquet...");

        // Get thread writers in order, and exactly as many as specified by self.num_files
        std::thread::scope(|scope| -> Result<()> {
            bvgraph
                .split_iter(num_files)
                .into_iter()
                .map(|mut arcs| {
                    let mut pl = pl.clone();
                    let fields = fields.clone();
                    let writer = &writer;
                    scope.spawn(move || {
                        let mut thread_writer = writer
                            .get_thread_writer()
                            .context("Could not get thread writer")?;
                        arcs.try_for_each(|(mut key, values)| {
                            let builder = thread_writer
                                .builder()
                                .context("Could not get table builder")?;
                            if builder.fields.is_none() {
                                builder.fields = Some(fields.clone());
                            }
                            if uncontract_keys {
                                key = contraction.underlying_node_id(key);
                            }

                            if key_swhids {
                                ensure_swhid_buffer(&mut builder.key_swhids);
                            }
                            if value_swhids {
                                ensure_swhid_buffer(&mut builder.value_swhids);
                            }

                            // write values
                            let mut num_values = 0;
                            values.into_iter().try_for_each(|mut value| -> Result<_> {
                                if uncontract_values {
                                    value = contraction.underlying_node_id(value);
                                }
                                builder.values.append_value(
                                    u64::try_from(value).context("Value overflowed u64")?,
                                );
                                if let Some(value_swhids) = builder.value_swhids.as_mut() {
                                    value_swhids
                                        .append_value(graph.properties().swhid(value).to_string());
                                }
                                num_values += 1;
                                Ok(())
                            })?;

                            // write keys
                            builder.keys.append_value_n(
                                u64::try_from(key).context("Key overflowed u64")?,
                                num_values,
                            );
                            if let Some(key_swhids) = builder.key_swhids.as_mut() {
                                for _ in 0..num_values {
                                    key_swhids
                                        .append_value(graph.properties().swhid(key).to_string());
                                }
                            }
                            pl.update_with_count(num_values);

                            Ok(())
                        })
                    })
                })
                .collect::<Vec<_>>()
                .into_iter()
                .try_for_each(|handle| handle.join().expect("thread panicked"))
        })?;

        pl.done();

        Ok(())
    }
}

fn ensure_swhid_buffer(buf: &mut Option<StringBuilder>) {
    if buf.is_none() {
        // FIXME: exclude validity buffer, values can't be null
        *buf = Some(StringBuilder::new());
    }
}

pub struct Builder {
    pub fields: Option<Vec<Field>>,
    pub keys: UInt64Builder,
    pub key_swhids: Option<StringBuilder>,
    pub values: UInt64Builder,
    pub value_swhids: Option<StringBuilder>,
}

impl Default for Builder {
    fn default() -> Self {
        Builder {
            fields: None,
            keys: UInt64Builder::new_from_buffer(
                Default::default(),
                None, // not nullable -> validity buffer not needed
            ),
            key_swhids: None,
            values: UInt64Builder::new_from_buffer(
                Default::default(),
                None, // not nullable -> validity buffer not needed
            ),
            value_swhids: None,
        }
    }
}

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

    fn buffer_size(&self) -> usize {
        // No validity slice
        self.len() * 8 // keys
        + self.len() * 8 // values
        + if self.key_swhids.is_some() { self.len() * (50 + 4) } else { 0 } // SWHID + offset
        + if self.value_swhids.is_some() { self.len() * (50 + 4) } else { 0 } // ditto
    }

    fn finish(&mut self) -> Result<StructArray> {
        let mut columns: Vec<Arc<dyn Array>> =
            vec![Arc::new(self.keys.finish()), Arc::new(self.values.finish())];
        if let Some(key_swhids) = self.key_swhids.as_mut() {
            columns.push(Arc::new(key_swhids.finish()));
        }
        if let Some(value_swhids) = self.value_swhids.as_mut() {
            columns.push(Arc::new(value_swhids.finish()));
        }

        let fields = self
            .fields
            .as_ref()
            .expect("bvgraph_to_parquet::Builder::fields not initialized");
        ensure!(
            columns.len() == fields.len(),
            "{} columns, but fields={fields:?}",
            columns.len()
        );

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