tenshift-core 0.1.3

//! Built-in sources for common data origins.
//!
//! - [`GlobSource`]  -  load files matching a glob pattern (e.g., `"data/*.json"`)
//! - [`MemorySource`]  -  load from an in-memory collection (for testing)
//! - [`ImageFolderSource`]  -  load files split into subdirectories by class label
//! - [`DistributedSampler`]  -  shard data across distributed training ranks
//! - [`JsonlSource`]  -  streaming JSON Lines parser
//! - `CsvSource`  -  load CSV files (feature: `csv`)
//! - `RingSource`  -  `io_uring`-accelerated batch file loading (feature: `uring`)
//!
//! # Creating Custom Sources
//!
//! Implement the [`Source`] trait to add new data origins:
//!
//! ```rust
//! use tenshift_core::source::{Source, SourceIterator};
//! use tenshift_core::sample::Sample;
//! use tenshift_core::error::Result;
//!
//! struct MySource;
//!
//! impl Source for MySource {
//!     fn open(&self) -> Result<Box<dyn SourceIterator>> {
//!         // Return an iterator over your data
//!         Ok(Box::new(std::iter::empty()))
//!     }
//!
//!     fn name(&self) -> &str { "my_source" }
//! }
//! ```

use std::path::{Path, PathBuf};

use crate::error::{Error, Result};
use crate::sample::{Sample, Tensor};
use crate::source::{Source, SourceIterator};

/// Collect every regular file matching a glob `pattern`, failing loud on any
/// per-path iteration error.
///
/// The former `glob::glob(pattern)?.filter_map(Result::ok)` silently discarded
/// each [`glob::GlobError`] (an I/O error, e.g. permission denied, while reading
/// a directory mid-iteration), so an unreadable subtree quietly shrank the
/// training set with no operator-visible signal - a Law-10 silent input loss in
/// the data pipeline. Each `GlobError` is now surfaced as [`Error::Io`]. Shared
/// by [`GlobSource::new`] and `RingSource::from_glob` (ONE-PLACE: previously the
/// identical filtering was hand-rolled in both).
pub(crate) fn collect_glob_files(pattern: &str) -> Result<Vec<PathBuf>> {
    let mut paths = Vec::new();
    for entry in glob::glob(pattern)? {
        let path = entry.map_err(|e| Error::Io(std::sync::Arc::new(e.into_error())))?;
        if path.is_file() {
            paths.push(path);
        }
    }
    Ok(paths)
}

#[cfg(feature = "image")]
pub mod image_folder;

#[cfg(feature = "image")]
pub use image_folder::ImageFolderSource;

#[cfg(feature = "csv")]
mod csv_source;
#[cfg(feature = "uring")]
mod ring_source;

#[cfg(feature = "csv")]
pub use csv_source::CsvSource;
#[cfg(feature = "uring")]
pub use ring_source::RingSource;

/// A source that reads files matching a glob pattern.
///
/// Each file becomes one sample with a `"data"` field containing the raw bytes
/// and metadata containing the filename.
///
/// ```rust
/// use tenshift_core::sources::GlobSource;
/// use tenshift_core::source::Source;
///
/// # fn example() -> tenshift_core::error::Result<()> {
/// let source = GlobSource::new("tests/data/*.json")?;
/// println!("found {} files", match source.len_hint() { Some(v) => v, None => 0 });
/// # Ok(())
/// # }
/// ```
pub struct GlobSource {
    paths: Vec<PathBuf>,
    pattern: String,
}

impl GlobSource {
    /// Create a source from a glob pattern.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidPattern`] if the pattern is invalid, or
    /// [`Error::EmptySource`] if no files match.
    pub fn new(pattern: &str) -> Result<Self> {
        let paths = collect_glob_files(pattern)?;

        if paths.is_empty() {
            return Err(Error::EmptySource {
                pattern: pattern.to_string(),
            });
        }

        tracing::info!(
            "GlobSource: found {} files matching '{}'",
            paths.len(),
            pattern
        );

        Ok(Self {
            paths,
            pattern: pattern.to_string(),
        })
    }

    /// Create a source from an explicit list of file paths.
    ///
    /// # Errors
    ///
    /// Returns [`Error::EmptySource`] if the path list is empty.
    pub fn from_paths(paths: Vec<PathBuf>) -> Result<Self> {
        if paths.is_empty() {
            return Err(Error::EmptySource {
                pattern: "<explicit paths>".to_string(),
            });
        }
        Ok(Self {
            pattern: format!("{} explicit paths", paths.len()),
            paths,
        })
    }
}

impl Source for GlobSource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        Ok(Box::new(GlobIterator {
            paths: self.paths.clone(),
            index: 0,
        }))
    }

    fn len_hint(&self) -> Option<u64> {
        Some(self.paths.len() as u64)
    }

    fn name(&self) -> &str {
        &self.pattern
    }
}

/// Iterator over files in a `GlobSource`.
struct GlobIterator {
    paths: Vec<PathBuf>,
    index: u64,
}

impl SourceIterator for GlobIterator {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        let Ok(idx) = usize::try_from(self.index) else {
            return Some(Err(Error::InvalidConfig {
                reason: "source index exceeded addressable memory on this platform".to_string(),
            }));
        };
        if idx >= self.paths.len() {
            return None;
        }

        let path = &self.paths[idx];
        let index = self.index;
        self.index += 1;

        Some(load_file_as_sample(path, index))
    }
}

/// Maximum JSONL line length (256MB). Prevents OOM from malicious/invalid
/// input files with unbounded lines.
const MAX_JSONL_LINE_LENGTH: usize = 256 * 1024 * 1024;

/// Load a single file into a Sample.
fn load_file_as_sample(path: &Path, index: u64) -> Result<Sample> {
    // Check size before reading to prevent OOM on huge files
    let metadata = std::fs::metadata(path).map_err(|e| Error::ReadFailed {
        path: path.to_path_buf(),
        reason: e.to_string(),
    })?;
    if metadata.len() > crate::pipeline::MAX_LOAD_FILE_SIZE {
        return Err(Error::ReadFailed {
            path: path.to_path_buf(),
            reason: format!(
                "file size {} exceeds maximum {} bytes",
                metadata.len(),
                crate::pipeline::MAX_LOAD_FILE_SIZE
            ),
        });
    }

    let data = std::fs::read(path).map_err(|e| Error::ReadFailed {
        path: path.to_path_buf(),
        reason: e.to_string(),
    })?;

    let filename = match path.file_name() {
        Some(n) => n.to_string_lossy().to_string(),
        None => String::new(),
    };

    Ok(Sample::new()
        .with("data", Tensor::bytes(data))
        .with("filename", Tensor::bytes(filename.as_bytes().to_vec()))
        .with_metadata(path.to_string_lossy(), index))
}

/// A source that yields samples from an in-memory collection.
///
/// Useful for testing and small datasets.
pub struct MemorySource {
    samples: Vec<Sample>,
    name: String,
}

impl MemorySource {
    /// Create a source from a vector of samples.
    pub fn new(name: impl Into<String>, samples: Vec<Sample>) -> Self {
        Self {
            samples,
            name: name.into(),
        }
    }
}

impl Source for MemorySource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        let samples: Vec<Result<Sample>> = self.samples.iter().cloned().map(Ok).collect();
        Ok(Box::new(samples.into_iter()))
    }

    fn len_hint(&self) -> Option<u64> {
        Some(self.samples.len() as u64)
    }

    fn name(&self) -> &str {
        &self.name
    }
}

/// A source wrapper that shards samples across distributed ranks.
///
/// Only samples where `index % world_size == rank` are yielded. The source
/// item's position within the underlying iterator determines the index.
///
/// ```rust
/// use tenshift_core::sample::Sample;
/// use tenshift_core::source::Source;
/// use tenshift_core::sources::{DistributedSampler, MemorySource};
///
/// let source = MemorySource::new("demo", vec![Sample::new(), Sample::new(), Sample::new()]);
/// let shard = DistributedSampler::new(source, 1, 2).unwrap();
/// assert_eq!(shard.len_hint(), Some(1));
/// ```
pub struct DistributedSampler<S> {
    inner: S,
    rank: usize,
    world_size: usize,
    name: String,
}

impl<S> DistributedSampler<S> {
    /// Create a distributed sampler over an existing source.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidConfig`] when `world_size == 0` or `rank >= world_size`.
    pub fn new(inner: S, rank: usize, world_size: usize) -> Result<Self>
    where
        S: Source,
    {
        validate_shard_config(rank, world_size)?;

        Ok(Self {
            name: format!("{}[rank={rank}/{world_size}]", inner.name()),
            inner,
            rank,
            world_size,
        })
    }
}

impl<S> Source for DistributedSampler<S>
where
    S: Source,
{
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        Ok(Box::new(DistributedSamplerIter {
            inner: self.inner.open()?,
            index: 0,
            rank: self.rank,
            world_size: self.world_size,
        }))
    }

    fn len_hint(&self) -> Option<u64> {
        self.inner
            .len_hint()
            .map(|len| distributed_len_hint(len, self.rank, self.world_size))
    }

    fn name(&self) -> &str {
        &self.name
    }
}

struct DistributedSamplerIter {
    inner: Box<dyn SourceIterator>,
    index: u64,
    rank: usize,
    world_size: usize,
}

impl SourceIterator for DistributedSamplerIter {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        loop {
            let item = self.inner.next_sample()?;
            let index = self.index;
            self.index = self.index.saturating_add(1);

            let in_shard = index % self.world_size as u64 == self.rank as u64;
            // Always forward errors regardless of shard - they represent source-level failures
            // that should be handled by ErrorPolicy, not silently dropped
            match item {
                Ok(sample) if in_shard => return Some(Ok(sample)),
                Err(error) => return Some(Err(error)),
                Ok(_) => {}
            }
        }
    }
}

fn validate_shard_config(rank: usize, world_size: usize) -> Result<()> {
    if world_size == 0 {
        return Err(Error::InvalidConfig {
            reason: "world_size must be greater than zero. Fix: pass a positive shard count."
                .to_string(),
        });
    }

    if rank >= world_size {
        return Err(Error::InvalidConfig {
            reason: format!(
                "rank {rank} is out of range for world_size {world_size}. Fix: use a rank in 0..{world_size}."
            ),
        });
    }

    Ok(())
}

fn distributed_len_hint(len: u64, rank: usize, world_size: usize) -> u64 {
    if len <= rank as u64 {
        0
    } else {
        ((len - 1 - rank as u64) / world_size as u64) + 1
    }
}

/// A JSONL (JSON Lines) source  -  one JSON object per line.
///
/// Each line becomes a sample with a `"json"` field containing the raw JSON bytes.
pub struct JsonlSource {
    path: PathBuf,
}

impl JsonlSource {
    /// Create a JSONL source from a file path.
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self { path: path.into() }
    }
}

impl Source for JsonlSource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        let file = std::fs::File::open(&self.path).map_err(|e| Error::ReadFailed {
            path: self.path.clone(),
            reason: e.to_string(),
        })?;
        let reader = std::io::BufReader::new(file);
        Ok(Box::new(JsonlIterator {
            reader,
            path: self.path.clone(),
            index: 0,
        }))
    }

    fn name(&self) -> &str {
        match self.path.to_str() {
            Some(v) => v,
            None => "jsonl",
        }
    }
}

/// Iterator over lines in a JSONL file.
struct JsonlIterator {
    reader: std::io::BufReader<std::fs::File>,
    path: PathBuf,
    index: u64,
}

impl SourceIterator for JsonlIterator {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        use std::io::BufRead;

        let mut line = String::new();
        match self.reader.read_line(&mut line) {
            Ok(0) => None, // EOF
            Ok(n) => {
                // Enforce maximum line length to prevent OOM from malicious input
                if n > MAX_JSONL_LINE_LENGTH {
                    return Some(Err(Error::CorruptData {
                        path: self.path.clone(),
                        reason: format!(
                            "line {} exceeds maximum length of {} bytes (got {} bytes). Fix: check for corrupt data or increase MAX_JSONL_LINE_LENGTH",
                            self.index + 1,
                            MAX_JSONL_LINE_LENGTH,
                            n
                        ),
                    }));
                }
                let trimmed = line.trim();
                if trimmed.is_empty() {
                    // Skip empty lines, recurse
                    return self.next_sample();
                }
                let index = self.index;
                self.index += 1;
                // Validate JSON syntax without building a DOM tree.
                // `RawValue` confirms the line is valid JSON but does not
                // allocate strings, maps, or vectors  -  just a bounds check.
                if let Err(error) = serde_json::from_str::<&serde_json::value::RawValue>(trimmed) {
                    return Some(Err(Error::CorruptData {
                        path: self.path.clone(),
                        reason: format!("invalid JSON on line {}: {error}", index + 1),
                    }));
                }
                Some(Ok(Sample::new()
                    .with("json", Tensor::bytes(trimmed.as_bytes().to_vec()))
                    .with_metadata(self.path.to_string_lossy(), index)))
            }
            Err(e) => Some(Err(Error::ReadFailed {
                path: self.path.clone(),
                reason: e.to_string(),
            })),
        }
    }
}

#[cfg(test)]
mod glob_tests {
    use super::collect_glob_files;

    #[test]
    fn collect_glob_files_returns_only_regular_files() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("a.bin"), "a").unwrap();
        std::fs::write(dir.path().join("b.bin"), "b").unwrap();
        // A subdirectory matching the glob must be excluded (not a regular file).
        std::fs::create_dir(dir.path().join("c.bin")).unwrap();
        // A non-matching file must be excluded by the pattern.
        std::fs::write(dir.path().join("d.txt"), "d").unwrap();

        let pattern = format!("{}/*.bin", dir.path().display());
        let mut files = collect_glob_files(&pattern).unwrap();
        files.sort();

        let names: Vec<_> = files
            .iter()
            .map(|p| p.file_name().unwrap().to_str().unwrap().to_string())
            .collect();
        assert_eq!(names, vec!["a.bin".to_string(), "b.bin".to_string()]);
    }

    #[test]
    fn collect_glob_files_rejects_invalid_pattern() {
        // An unclosed `[` is an invalid glob pattern; it must surface as an
        // error, not silently return an empty file list.
        let err = collect_glob_files("data/[.bin");
        assert!(err.is_err(), "invalid glob pattern must error, got {err:?}");
    }
}