// 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::Result;
use crossbeam::atomic::AtomicCell;
use dsi_progress_logger::{concurrent_progress_logger, ProgressLog};
use rayon::prelude::*;
use rdst::RadixSort;
use sux::bits::BitFieldVec;
use sux::func::{VBuilder, VFunc};
use sux::prelude::BitFieldSlice;
use sux::utils::FromIntoIterator;

use crate::table::Table;
use crate::{Sha1, Sha1Git, SHA1_FILENAME, SHA1_GIT_FILENAME};

#[cfg(feature = "orc")]
mod orc;
#[cfg(feature = "orc")]
pub use orc::*;

pub fn build_sha1_git_table_from_keys(
    mut sha1_gits: Vec<Sha1Git>,
    num_rows: usize,
    tables_dir: &Path,
    parallelism: NonZeroUsize,
) -> Result<Table<Sha1Git, impl AsRef<[u8]>>> {
    let mut pl = concurrent_progress_logger! {
        item_name="row",
        display_memory=true,
        expected_updates=Some(num_rows),
    };

    pl.start("Listing keys...");

    pl.done();

    log::info!("Sorting keys...");
    sha1_gits.radix_sort_unstable();

    log::info!("Deduplicating sha1_gits...");
    sha1_gits.dedup();
    let num_duplicates = num_rows - sha1_gits.len();
    if num_duplicates > 0 {
        log::warn!("Removed {num_duplicates} duplicates");
    } else {
        log::info!("No duplicates found");
    }

    log::info!("Writing sha1_git table...");
    Table::from_array(tables_dir.join(SHA1_GIT_FILENAME), &sha1_gits, parallelism)
}

pub fn build_vfunc(
    keys: &Table<Sha1Git, impl AsRef<[u8]>>,
) -> Result<VFunc<Sha1Git, usize, BitFieldVec>> {
    let num_keys = keys.len();

    let mut pl = concurrent_progress_logger! {
        item_name="key",
        display_memory=true,
        expected_updates=Some(num_keys),
    };
    pl.start("Building VFunc...");
    let builder = VBuilder::<_, BitFieldVec<usize>>::default().expected_num_keys(num_keys);
    let func = builder.try_build_func(
        FromIntoIterator::from(keys.into_iter()),
        FromIntoIterator::from(0..num_keys),
        &mut pl,
    )?;
    pl.done();

    Ok(func)
}

pub fn build_sha1_table_from_pairs<I, D>(
    pairs_iter: I,
    vfunc: &VFunc<Sha1Git, usize, D>,
    num_keys: usize,
    tables_dir: &Path,
    parallelism: NonZeroUsize,
) -> Result<()>
where
    I: ParallelIterator<Item = Result<(Sha1Git, Sha1)>>,
    D: BitFieldSlice<usize> + Sync,
{
    log::info!("Allocating arrays...");
    let sha1s: Vec<_> = (0..num_keys)
        .into_par_iter()
        .map(|_| AtomicCell::default())
        .collect();

    let mut pl = concurrent_progress_logger!(
        item_name = "row",
        display_memory = true,
        expected_updates = Some(num_keys),
    );

    pl.start("Building tables...");

    let res = pairs_iter.try_for_each_with(pl.clone(), |pl, pair| {
        let (sha1_git, sha1) = pair?;
        let index = vfunc.get(sha1_git);
        sha1s[index].store(sha1);
        pl.light_update();
        Ok(())
    });

    pl.done();

    // Compiled to no-op
    let sha1s: Vec<_> = sha1s.into_iter().map(AtomicCell::into_inner).collect();

    log::info!("Writing sha1 table...");
    Table::from_array(tables_dir.join(SHA1_FILENAME), &sha1s, parallelism)?;

    res
}
