// 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::borrow::Borrow;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
#[cfg(feature = "arrow")]
use dataset_writer::{
    parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder,
    parquet::file::reader::ChunkReader,
};
use sux::traits::{IndexedDict, IndexedSeq, Types};

pub fn list_dir(dir_path: &Path) -> Result<Vec<PathBuf>> {
    Ok(std::fs::read_dir(dir_path)
        .with_context(|| format!("Could not stat directory {}", dir_path.display()))?
        .map(|entry| Ok(entry?))
        .collect::<Result<Vec<_>>>()
        .with_context(|| format!("Could not stat {} entry", dir_path.display()))?
        .into_iter()
        .map(|entry| entry.path().to_path_buf())
        .collect())
}

#[cfg(feature = "arrow")]
pub fn open_parquet_table(
    dir_path: &Path,
) -> Result<Vec<ParquetRecordBatchReaderBuilder<impl ChunkReader + use<>>>> {
    use rayon::prelude::*;
    use std::fs::File;

    list_dir(dir_path)?
        .into_par_iter()
        .flat_map(|path| {
            if path.extension() != Some(std::ffi::OsStr::new("parquet")) {
                log::info!("Skipping {}, not .parquet", path.display());
                return None;
            }
            Some(
                File::open(&path)
                    .with_context(|| format!("Could not open {}", path.display()))
                    .and_then(|file| {
                        ParquetRecordBatchReaderBuilder::try_new(file).with_context(|| {
                            format!("Could not open {} as Parquet", path.display())
                        })
                    }),
            )
        })
        .collect::<Result<Vec<_>>>()
}

/// An implementation of [`ContractionBackend`](swh_graph::views::contiguous_subgraph::ContractionBackend)
/// that does not actually renumber nodes
pub struct IdentityContractionBackend {
    pub num_nodes: usize,
}

impl Types for IdentityContractionBackend {
    type Input = usize;
    type Output<'a> = usize;
}
impl IndexedSeq for IdentityContractionBackend {
    unsafe fn get_unchecked(&self, n: usize) -> usize {
        assert!(n < self.num_nodes);
        n
    }

    fn len(&self) -> usize {
        self.num_nodes
    }
}
impl IndexedDict for IdentityContractionBackend {
    fn index_of(&self, n: impl Borrow<usize>) -> Option<usize> {
        let n = *n.borrow();
        assert!(n < self.num_nodes);
        Some(n)
    }
}
