// Copyright (C) 2025-2026  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::num::NonZeroUsize;
use std::path::Path;

use anyhow::{Context, Result};
use ar_row::deserialize::ArRowStruct;
use ar_row::row_iterator::RowIterator;
use ar_row_derive::ArRowDeserialize;
use dsi_progress_logger::{concurrent_progress_logger, ProgressLog};
use epserde::deser::{Deserialize, Flags};
use epserde::ser::Serialize;
use orc_rust::projection::ProjectionMask;
use orc_rust::ArrowReaderBuilder;
use rayon::prelude::*;
use sux::func::VFunc;
use sux::prelude::BitFieldSlice;

use crate::table::Table;
use crate::{Sha1, Sha1Git, VFUNC_FILENAME};

use super::{build_sha1_git_table_from_keys, build_sha1_table_from_pairs, build_vfunc};

pub(crate) fn get_dataset_reader_builders<P: AsRef<Path>>(
    dataset_dir: P,
    subdirectory: &str,
) -> Result<Vec<ArrowReaderBuilder<std::fs::File>>> {
    let mut dataset_dir = dataset_dir.as_ref().to_owned();
    dataset_dir.push(subdirectory);
    std::fs::read_dir(&dataset_dir)
        .with_context(|| format!("Could not list {}", dataset_dir.display()))?
        .map(|file_path| {
            let file_path = file_path
                .with_context(|| format!("Failed to list {}", dataset_dir.display()))?
                .path();
            let file = std::fs::File::open(&file_path)
                .with_context(|| format!("Could not open {}", file_path.display()))?;
            let builder = ArrowReaderBuilder::try_new(file)
                .with_context(|| format!("Could not read {}", file_path.display()))?;
            Ok(builder)
        })
        .collect()
}

#[deprecated(
    since = "0.1.5",
    note = "please use `build_sha1_git_table_from_orc_export` instead"
)]
/// Deprecated alias for [`build_sha1_git_table_from_orc_export`]
pub fn build_sha1_git_table<P: AsRef<Path>>(
    orc_dir: P,
    tables_dir: P,
    parallelism: NonZeroUsize,
) -> Result<Table<Sha1Git, impl AsRef<[u8]>>> {
    build_sha1_git_table_from_orc_export(orc_dir, tables_dir, parallelism)
}

pub fn build_sha1_git_table_from_orc_export<P: AsRef<Path>>(
    orc_dir: P,
    tables_dir: P,
    parallelism: NonZeroUsize,
) -> Result<Table<Sha1Git, impl AsRef<[u8]>>> {
    let tables_dir = tables_dir.as_ref();
    let dataset_reader_builders = get_dataset_reader_builders(orc_dir, "content")?;

    let num_rows = usize::try_from(
        dataset_reader_builders
            .iter()
            .map(|reader_builder| reader_builder.file_metadata().number_of_rows())
            .sum::<u64>(),
    )
    .context("Number of rows overflowed usize")?;

    let mut pl = concurrent_progress_logger! {
        item_name="row",
        display_memory=true,
        expected_updates=Some(num_rows),
    };
    pl.start("Building sha1_git table...");

    let sha1_gits = dataset_reader_builders
        .into_par_iter()
        .flat_map_iter(move |reader_builder| {
            #[derive(ArRowDeserialize, Clone, Default)]
            struct Row {
                sha1_git: String,
            }

            let mut pl = pl.clone();

            let projection = ProjectionMask::named_roots(
                reader_builder.file_metadata().root_data_type(),
                &Row::columns(),
            );
            let reader = reader_builder.with_projection(projection).build();
            RowIterator::new(reader.map(|batch| batch.unwrap()))
                .expect("Could not create iterator") // TODO: don't panic
                .map(move |Row { sha1_git }| {
                    pl.light_update();
                    Sha1Git::from_hex(sha1_git.as_bytes())
                })
        })
        .collect::<Result<_>>()?;

    build_sha1_git_table_from_keys(sha1_gits, num_rows, tables_dir, parallelism)
}

fn build_sha1_table_from_orc_export<I, D>(
    reader_builders_iter: I,
    vfunc: &VFunc<Sha1Git, usize, D>,
    num_keys: usize,
    tables_dir: &Path,
    parallelism: NonZeroUsize,
) -> Result<()>
where
    I: IntoParallelIterator<Item = ArrowReaderBuilder<std::fs::File>>,
    D: BitFieldSlice<usize> + Sync,
{
    let pairs_iter = reader_builders_iter
        .into_par_iter()
        .flat_map_iter(|reader_builder| {
            #[derive(ArRowDeserialize, Clone, Default)]
            struct Row {
                sha1: String,
                sha1_git: String,
            }

            let projection = ProjectionMask::named_roots(
                reader_builder.file_metadata().root_data_type(),
                &Row::columns(),
            );
            let reader = reader_builder.with_projection(projection).build();
            RowIterator::new(reader.map(|batch| batch.unwrap()))
                .expect("Could not create iterator") // TODO: don't panic
                .map(|Row { sha1_git, sha1 }| {
                    let sha1 = Sha1::from_hex(sha1.as_bytes())?;
                    let sha1_git = Sha1Git::from_hex(sha1_git.as_bytes())?;
                    Ok((sha1_git, sha1))
                })
        });

    build_sha1_table_from_pairs(pairs_iter, vfunc, num_keys, tables_dir, parallelism)
}

#[deprecated(
    since = "0.1.5",
    note = "please use `build_tables_from_orc_export` instead"
)]
/// Deprecated alias for [`build_tables_from_orc_export`]
pub fn build_tables<P: AsRef<Path>, D: BitFieldSlice<usize> + Sync>(
    orc_dir: P,
    vfunc: &VFunc<Sha1Git, usize, D>,
    tables_dir: P,
    parallelism: NonZeroUsize,
) -> Result<()> {
    build_tables_from_orc_export(orc_dir, vfunc, tables_dir, parallelism)
}

pub fn build_tables_from_orc_export<P: AsRef<Path>, D: BitFieldSlice<usize> + Sync>(
    orc_dir: P,
    vfunc: &VFunc<Sha1Git, usize, D>,
    tables_dir: P,
    parallelism: NonZeroUsize,
) -> Result<()> {
    let tables_dir = tables_dir.as_ref();
    let num_keys = vfunc.len();
    let reader_builders = get_dataset_reader_builders(orc_dir, "content")?;

    build_sha1_table_from_orc_export(reader_builders, vfunc, num_keys, tables_dir, parallelism)
}

#[deprecated(
    since = "0.1.5",
    note = "please use `build_digestmap_from_orc_export` instead"
)]
/// Deprecated alias for [`build_digestmap_from_orc_export`]
pub fn build_digestmap(orc: &Path, dir_out: &Path) -> Result<()> {
    build_digestmap_from_orc_export(orc, dir_out)
}

pub fn build_digestmap_from_orc_export(orc: &Path, dir_out: &Path) -> Result<()> {
    std::fs::create_dir_all(dir_out)
        .with_context(|| format!("Could not create {}", dir_out.display()))?;

    let parallelism = NonZeroUsize::new(8).unwrap(); // TODO: configurable

    let vfunc_path = dir_out.join(VFUNC_FILENAME);
    let load_vfunc = std::fs::exists(&vfunc_path)
        .with_context(|| format!("Could not check if {} exists", vfunc_path.display()))?;
    if load_vfunc {
        log::info!("Load vfunc from {}", vfunc_path.display());
        let vfunc = <VFunc<_>>::mmap(
            &vfunc_path,
            Flags::TRANSPARENT_HUGE_PAGES | Flags::RANDOM_ACCESS,
        )?;
        build_tables_from_orc_export(orc, &*vfunc, dir_out, parallelism)
            .context("Could not build tables")?;
    } else {
        let sha1_gits = build_sha1_git_table_from_orc_export(orc, dir_out, parallelism)
            .context("Could not build sha1_git table")?;
        let vfunc = build_vfunc(&sha1_gits).context("Could not build VFunc")?;

        log::info!("Writing VFunc...");
        vfunc
            .store(&vfunc_path)
            .with_context(|| format!("Could not save VFunc to {}", vfunc_path.display()))?;
        build_tables_from_orc_export(orc, &vfunc, dir_out, parallelism)
            .context("Could not build tables")?;
    };

    Ok(())
}
