provisioning-platform/crates/rag/src/ingestion.rs

//! Document ingestion pipeline

use std::path::Path;

use walkdir::WalkDir;

use crate::chunking::ChunkingEngine;
use crate::config::IngestionConfig;
use crate::embeddings::{EmbeddedDocument, EmbeddingEngine};
use crate::error::Result;

/// Document ingester for RAG system
pub struct DocumentIngester {
    config: IngestionConfig,
    embedding_engine: EmbeddingEngine,
    chunking_engine: ChunkingEngine,
}

impl DocumentIngester {
    /// Create a new document ingester
    pub fn new(config: IngestionConfig, embedding_engine: EmbeddingEngine) -> Self {
        let chunking_engine = ChunkingEngine::new(config.chunk_size, config.chunk_overlap);

        Self {
            config,
            embedding_engine,
            chunking_engine,
        }
    }

    /// Ingest a single document file
    pub async fn ingest_file(&self, file_path: &Path) -> Result<Vec<EmbeddedDocument>> {
        let path_str = file_path.to_string_lossy().to_string();

        // Read file content
        let content = std::fs::read_to_string(file_path).map_err(|e| {
            tracing::error!("Failed to read file {}: {}", path_str, e);
            crate::error::RagError::Io(e)
        })?;

        // Chunk the document
        let chunks = self.chunking_engine.chunk_file(&content, &path_str)?;

        tracing::debug!("Chunked {} into {} chunks", path_str, chunks.len());

        // Embed the chunks
        let embedded = self.embedding_engine.embed_batch(&chunks).await?;

        tracing::info!(
            "Ingested {} with {} embedded chunks",
            path_str,
            embedded.len()
        );

        Ok(embedded)
    }

    /// Ingest markdown documentation
    pub async fn ingest_markdown(&self, doc_path: &Path) -> Result<Vec<EmbeddedDocument>> {
        let content = std::fs::read_to_string(doc_path)?;
        let path_str = doc_path.to_string_lossy().to_string();

        let chunks = self.chunking_engine.chunk_markdown(&content, &path_str)?;
        let embedded = self.embedding_engine.embed_batch(&chunks).await?;

        tracing::info!(
            "Ingested markdown document: {} ({} chunks)",
            path_str,
            embedded.len()
        );

        Ok(embedded)
    }

    /// Ingest KCL schema files
    pub async fn ingest_kcl(&self, schema_path: &Path) -> Result<Vec<EmbeddedDocument>> {
        let content = std::fs::read_to_string(schema_path)?;
        let path_str = schema_path.to_string_lossy().to_string();

        let chunks = self.chunking_engine.chunk_kcl(&content, &path_str)?;
        let embedded = self.embedding_engine.embed_batch(&chunks).await?;

        tracing::info!(
            "Ingested KCL schema: {} ({} chunks)",
            path_str,
            embedded.len()
        );

        Ok(embedded)
    }

    /// Ingest Nushell scripts
    pub async fn ingest_nushell(&self, script_path: &Path) -> Result<Vec<EmbeddedDocument>> {
        let content = std::fs::read_to_string(script_path)?;
        let path_str = script_path.to_string_lossy().to_string();

        let chunks = self.chunking_engine.chunk_nushell(&content, &path_str)?;
        let embedded = self.embedding_engine.embed_batch(&chunks).await?;

        tracing::info!(
            "Ingested Nushell script: {} ({} chunks)",
            path_str,
            embedded.len()
        );

        Ok(embedded)
    }

    /// Scan directory recursively for documents to ingest
    pub async fn scan_and_ingest_directory(
        &self,
        root_path: &Path,
    ) -> Result<Vec<EmbeddedDocument>> {
        let mut all_documents = Vec::new();

        for entry in WalkDir::new(root_path)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.path().is_file())
        {
            let path = entry.path();
            let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");

            // Check if this file type should be ingested
            let should_ingest = self.config.doc_types.iter().any(|t| {
                (t == "markdown" && ext == "md")
                    || (t == "kcl" && ext == "k")
                    || (t == "nushell" && ext == "nu")
            });

            if should_ingest {
                match self.ingest_file(path).await {
                    Ok(docs) => {
                        all_documents.extend(docs);
                    }
                    Err(e) => {
                        tracing::warn!("Failed to ingest {}: {}", path.display(), e);
                        // Continue with next file
                    }
                }
            }
        }

        tracing::info!(
            "Ingested {} total documents from {}",
            all_documents.len(),
            root_path.display()
        );

        Ok(all_documents)
    }

    /// Scan directory for documents without ingesting
    pub async fn scan_documents(&self, root_path: &Path) -> Result<Vec<String>> {
        let mut documents = Vec::new();

        for entry in WalkDir::new(root_path)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.path().is_file())
        {
            let path = entry.path();
            let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");

            if self.config.doc_types.iter().any(|t| {
                (t == "markdown" && ext == "md")
                    || (t == "kcl" && ext == "k")
                    || (t == "nushell" && ext == "nu")
            }) {
                documents.push(path.to_string_lossy().to_string());
            }
        }

        tracing::debug!("Found {} documents to ingest", documents.len());

        Ok(documents)
    }

    /// Get reference to embedding engine
    pub fn embedding_engine(&self) -> &EmbeddingEngine {
        &self.embedding_engine
    }

    /// Get reference to chunking engine
    pub fn chunking_engine(&self) -> &ChunkingEngine {
        &self.chunking_engine
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::EmbeddingConfig;

    #[test]
    fn test_document_ingester_creation() {
        let embedding_config = EmbeddingConfig {
            provider: "local".to_string(),
            ..Default::default()
        };

        let engine = EmbeddingEngine::new(embedding_config).unwrap();
        let ingest_config = IngestionConfig::default();

        let _ingester = DocumentIngester::new(ingest_config, engine);
        // Basic smoke test - object created successfully
    }

    #[tokio::test]
    async fn test_scan_documents() {
        let embedding_config = EmbeddingConfig {
            provider: "local".to_string(),
            ..Default::default()
        };

        let engine = EmbeddingEngine::new(embedding_config).unwrap();
        let ingest_config = IngestionConfig::default();

        let ingester = DocumentIngester::new(ingest_config, engine);

        // Try to scan a directory that doesn't exist
        // This should return empty list, not error
        let result = ingester.scan_documents(Path::new("/nonexistent")).await;

        // Should either succeed with empty list or fail gracefully
        match result {
            Ok(docs) => {
                assert!(docs.is_empty());
            }
            Err(_) => {
                // Also acceptable - directory doesn't exist
            }
        }
    }
}