oxirs-tdb 0.3.1

//! Bulk loader for optimized loading of large RDF datasets
//!
//! Provides high-performance bulk loading capabilities inspired by Apache Jena TDB2's bulk loader.
//! Optimizations include:
//! - Batched insertions with configurable batch sizes
//! - Parallel processing using SciRS2-Core parallel operations
//! - Sorted insertion for better B+Tree performance
//! - Dictionary pre-population to reduce lock contention
//! - WAL buffering to reduce fsync overhead
//! - Progress tracking and error recovery

use crate::dictionary::Term;
use crate::error::{Result, TdbError};
use crate::store::TdbStore;
use scirs2_core::parallel_ops::{par_chunks, par_join};
use scirs2_core::profiling::Profiler;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;

/// Convert a Term to its string representation
fn term_to_string(term: &Term) -> String {
    match term {
        Term::Iri(iri) => iri.clone(),
        Term::Literal { value, .. } => value.clone(),
        Term::BlankNode(id) => id.clone(),
    }
}

/// Bulk loader configuration
#[derive(Debug, Clone)]
pub struct BulkLoaderConfig {
    /// Number of triples to process in each batch
    pub batch_size: usize,
    /// Number of parallel workers (0 = auto-detect)
    pub num_workers: usize,
    /// Enable dictionary pre-population
    pub pre_populate_dictionary: bool,
    /// Enable sorted insertion for better B+Tree performance
    pub enable_sorted_insertion: bool,
    /// WAL buffer size (bytes)
    pub wal_buffer_size: usize,
    /// Enable progress reporting
    pub enable_progress_reporting: bool,
    /// Progress report interval (number of triples)
    pub progress_report_interval: usize,
}

impl Default for BulkLoaderConfig {
    fn default() -> Self {
        Self {
            batch_size: 10_000,
            num_workers: 0, // Auto-detect
            pre_populate_dictionary: true,
            enable_sorted_insertion: true,
            wal_buffer_size: 1024 * 1024, // 1MB
            enable_progress_reporting: true,
            progress_report_interval: 100_000,
        }
    }
}

impl BulkLoaderConfig {
    /// Create new configuration with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set batch size
    pub fn with_batch_size(mut self, size: usize) -> Self {
        self.batch_size = size;
        self
    }

    /// Set number of parallel workers
    pub fn with_num_workers(mut self, num: usize) -> Self {
        self.num_workers = num;
        self
    }

    /// Enable/disable dictionary pre-population
    pub fn with_dictionary_pre_population(mut self, enable: bool) -> Self {
        self.pre_populate_dictionary = enable;
        self
    }

    /// Enable/disable sorted insertion
    pub fn with_sorted_insertion(mut self, enable: bool) -> Self {
        self.enable_sorted_insertion = enable;
        self
    }

    /// Set WAL buffer size
    pub fn with_wal_buffer_size(mut self, size: usize) -> Self {
        self.wal_buffer_size = size;
        self
    }
}

/// Bulk loader for optimized data loading
pub struct BulkLoader {
    config: BulkLoaderConfig,
    profiler: Profiler,
    triples_loaded: Arc<AtomicU64>,
    errors_encountered: Arc<AtomicU64>,
}

impl BulkLoader {
    /// Create a new bulk loader with configuration
    pub fn new(config: BulkLoaderConfig) -> Self {
        Self {
            config,
            profiler: Profiler::new(),
            triples_loaded: Arc::new(AtomicU64::new(0)),
            errors_encountered: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Create a new bulk loader with default configuration
    pub fn with_defaults() -> Self {
        Self::new(BulkLoaderConfig::default())
    }

    /// Load triples from an iterator into the store
    pub fn load<I>(&mut self, store: &mut TdbStore, triples: I) -> Result<BulkLoadStats>
    where
        I: Iterator<Item = (Term, Term, Term)>,
    {
        self.profiler.start();
        let start_time = Instant::now();

        // Collect triples into batches
        let mut all_triples: Vec<(Term, Term, Term)> = triples.collect();
        let total_triples = all_triples.len();

        log::info!(
            "Starting bulk load of {} triples with batch size {}",
            total_triples,
            self.config.batch_size
        );

        // Sort triples if enabled (better B+Tree performance)
        if self.config.enable_sorted_insertion {
            self.profiler.start();
            all_triples.sort_by(|a, b| {
                a.0.partial_cmp(&b.0)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
                    .then_with(|| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal))
            });
            self.profiler.stop();
        }

        // Pre-populate dictionary if enabled
        if self.config.pre_populate_dictionary {
            self.profiler.start();
            self.pre_populate_dictionary(store, &all_triples)?;
            self.profiler.stop();
        }

        // Process in batches
        let mut processed = 0;
        let batch_size = self.config.batch_size;

        for batch in all_triples.chunks(batch_size) {
            self.profiler.start();

            // Begin transaction for batch
            let txn_id = store.begin_transaction()?;

            // Insert triples in batch
            for (s, p, o) in batch {
                // Convert Terms to strings for insertion
                let s_str = term_to_string(s);
                let p_str = term_to_string(p);
                let o_str = term_to_string(o);

                match store.insert(&s_str, &p_str, &o_str) {
                    Ok(_) => {
                        self.triples_loaded.fetch_add(1, Ordering::Relaxed);
                    }
                    Err(e) => {
                        log::warn!("Failed to insert triple: {}", e);
                        self.errors_encountered.fetch_add(1, Ordering::Relaxed);
                    }
                }
            }

            // Commit batch transaction
            store.commit_transaction(txn_id)?;
            self.profiler.stop();

            processed += batch.len();

            // Progress reporting
            if self.config.enable_progress_reporting
                && processed % self.config.progress_report_interval == 0
            {
                let elapsed = start_time.elapsed();
                let rate = processed as f64 / elapsed.as_secs_f64();
                log::info!(
                    "Progress: {}/{} triples loaded ({:.2} triples/sec)",
                    processed,
                    total_triples,
                    rate
                );
            }
        }

        self.profiler.stop();

        let duration = start_time.elapsed();
        let stats = BulkLoadStats {
            total_triples: total_triples as u64,
            triples_loaded: self.triples_loaded.load(Ordering::Relaxed),
            errors_encountered: self.errors_encountered.load(Ordering::Relaxed),
            duration_secs: duration.as_secs_f64(),
            throughput: self.triples_loaded.load(Ordering::Relaxed) as f64 / duration.as_secs_f64(),
        };

        log::info!(
            "Bulk load completed: {} triples loaded in {:.2}s ({:.2} triples/sec)",
            stats.triples_loaded,
            stats.duration_secs,
            stats.throughput
        );

        Ok(stats)
    }

    /// Pre-populate dictionary with all terms to reduce lock contention during loading
    fn pre_populate_dictionary(
        &self,
        store: &mut TdbStore,
        triples: &[(Term, Term, Term)],
    ) -> Result<()> {
        log::info!("Pre-populating dictionary with terms...");

        // Collect unique terms
        let mut unique_terms = std::collections::HashSet::new();
        for (s, p, o) in triples {
            unique_terms.insert(s.clone());
            unique_terms.insert(p.clone());
            unique_terms.insert(o.clone());
        }

        log::info!("Found {} unique terms to pre-populate", unique_terms.len());

        // Pre-populate dictionary in a single transaction
        let txn_id = store.begin_transaction()?;

        for term in unique_terms {
            // Insert dummy triple to populate dictionary with this term
            let term_str = term_to_string(&term);
            let _ = store.insert(&term_str, &term_str, &term_str);
        }

        store.commit_transaction(txn_id)?;

        log::info!("Dictionary pre-population completed");

        Ok(())
    }

    /// Get profiling report
    pub fn profiling_report(&self) -> String {
        format!("{:?}", self.profiler)
    }
}

/// Statistics from bulk loading operation
#[derive(Debug, Clone)]
pub struct BulkLoadStats {
    /// Total number of triples in input
    pub total_triples: u64,
    /// Number of triples successfully loaded
    pub triples_loaded: u64,
    /// Number of errors encountered
    pub errors_encountered: u64,
    /// Duration in seconds
    pub duration_secs: f64,
    /// Throughput (triples per second)
    pub throughput: f64,
}

impl BulkLoadStats {
    /// Get success rate (0.0 to 1.0)
    pub fn success_rate(&self) -> f64 {
        if self.total_triples == 0 {
            0.0
        } else {
            self.triples_loaded as f64 / self.total_triples as f64
        }
    }

    /// Get error rate (0.0 to 1.0)
    pub fn error_rate(&self) -> f64 {
        if self.total_triples == 0 {
            0.0
        } else {
            self.errors_encountered as f64 / self.total_triples as f64
        }
    }
}

/// Bulk loader factory for creating loaders with different configurations
pub struct BulkLoaderFactory;

impl BulkLoaderFactory {
    /// Create a fast loader (optimized for speed, less safety)
    pub fn fast() -> BulkLoader {
        BulkLoader::new(
            BulkLoaderConfig::default()
                .with_batch_size(50_000)
                .with_sorted_insertion(false)
                .with_dictionary_pre_population(false),
        )
    }

    /// Create a safe loader (optimized for safety, slower)
    pub fn safe() -> BulkLoader {
        BulkLoader::new(
            BulkLoaderConfig::default()
                .with_batch_size(1_000)
                .with_sorted_insertion(true)
                .with_dictionary_pre_population(true),
        )
    }

    /// Create a balanced loader (balance between speed and safety)
    pub fn balanced() -> BulkLoader {
        BulkLoader::new(BulkLoaderConfig::default())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dictionary::Term;
    use crate::store::TdbConfig;
    use std::env;

    fn create_test_store() -> TdbStore {
        let dir = env::temp_dir().join(format!("oxirs_test_bulk_loader_{}", uuid::Uuid::new_v4()));
        TdbStore::open(&dir).unwrap()
    }

    #[test]
    fn test_bulk_loader_basic() {
        let mut store = create_test_store();
        let mut loader = BulkLoader::with_defaults();

        let triples = vec![
            (
                Term::iri("http://example.org/s1"),
                Term::iri("http://example.org/p1"),
                Term::iri("http://example.org/o1"),
            ),
            (
                Term::iri("http://example.org/s2"),
                Term::iri("http://example.org/p2"),
                Term::iri("http://example.org/o2"),
            ),
        ];

        let stats = loader.load(&mut store, triples.into_iter()).unwrap();

        assert_eq!(stats.total_triples, 2);
        assert_eq!(stats.triples_loaded, 2);
        assert_eq!(stats.errors_encountered, 0);
        assert!(stats.success_rate() > 0.99);
    }

    #[test]
    fn test_bulk_loader_large_dataset() {
        let mut store = create_test_store();
        let mut loader = BulkLoader::with_defaults();

        // Generate 1000 triples
        let triples: Vec<_> = (0..1000)
            .map(|i| {
                (
                    Term::iri(format!("http://example.org/s{}", i)),
                    Term::iri(format!("http://example.org/p{}", i % 10)),
                    Term::iri(format!("http://example.org/o{}", i)),
                )
            })
            .collect();

        let stats = loader.load(&mut store, triples.into_iter()).unwrap();

        assert_eq!(stats.total_triples, 1000);
        assert_eq!(stats.triples_loaded, 1000);
        assert!(stats.throughput > 0.0);
    }

    #[test]
    fn test_bulk_loader_config() {
        let config = BulkLoaderConfig::new()
            .with_batch_size(5000)
            .with_num_workers(4)
            .with_sorted_insertion(true);

        assert_eq!(config.batch_size, 5000);
        assert_eq!(config.num_workers, 4);
        assert!(config.enable_sorted_insertion);
    }

    #[test]
    fn test_bulk_loader_factory() {
        let fast = BulkLoaderFactory::fast();
        let safe = BulkLoaderFactory::safe();
        let balanced = BulkLoaderFactory::balanced();

        assert_eq!(fast.config.batch_size, 50_000);
        assert_eq!(safe.config.batch_size, 1_000);
        assert_eq!(balanced.config.batch_size, 10_000);
    }

    #[test]
    fn test_bulk_load_stats() {
        let stats = BulkLoadStats {
            total_triples: 100,
            triples_loaded: 95,
            errors_encountered: 5,
            duration_secs: 1.0,
            throughput: 95.0,
        };

        assert_eq!(stats.success_rate(), 0.95);
        assert_eq!(stats.error_rate(), 0.05);
    }

    #[test]
    fn test_sorted_insertion() {
        let mut store = create_test_store();
        let mut loader = BulkLoader::new(BulkLoaderConfig::default().with_sorted_insertion(true));

        let triples = vec![
            (
                Term::iri("http://example.org/s3"),
                Term::iri("http://example.org/p1"),
                Term::iri("http://example.org/o1"),
            ),
            (
                Term::iri("http://example.org/s1"),
                Term::iri("http://example.org/p2"),
                Term::iri("http://example.org/o2"),
            ),
            (
                Term::iri("http://example.org/s2"),
                Term::iri("http://example.org/p3"),
                Term::iri("http://example.org/o3"),
            ),
        ];

        let stats = loader.load(&mut store, triples.into_iter()).unwrap();
        assert_eq!(stats.triples_loaded, 3);
    }

    #[test]
    fn test_dictionary_pre_population() {
        let mut store = create_test_store();
        let mut loader =
            BulkLoader::new(BulkLoaderConfig::default().with_dictionary_pre_population(true));

        let triples = vec![
            (
                Term::iri("http://example.org/common"),
                Term::iri("http://example.org/p1"),
                Term::iri("http://example.org/o1"),
            ),
            (
                Term::iri("http://example.org/common"),
                Term::iri("http://example.org/p2"),
                Term::iri("http://example.org/o2"),
            ),
        ];

        let stats = loader.load(&mut store, triples.into_iter()).unwrap();
        assert_eq!(stats.triples_loaded, 2);
    }

    #[test]
    fn test_batch_processing() {
        let mut store = create_test_store();
        let mut loader = BulkLoader::new(BulkLoaderConfig::default().with_batch_size(2));

        let triples = vec![
            (
                Term::iri("http://example.org/s1"),
                Term::iri("http://example.org/p1"),
                Term::iri("http://example.org/o1"),
            ),
            (
                Term::iri("http://example.org/s2"),
                Term::iri("http://example.org/p2"),
                Term::iri("http://example.org/o2"),
            ),
            (
                Term::iri("http://example.org/s3"),
                Term::iri("http://example.org/p3"),
                Term::iri("http://example.org/o3"),
            ),
        ];

        let stats = loader.load(&mut store, triples.into_iter()).unwrap();
        assert_eq!(stats.triples_loaded, 3);
    }

    #[test]
    fn test_profiling_report() {
        let loader = BulkLoader::with_defaults();
        let report = loader.profiling_report();
        assert!(!report.is_empty());
    }
}