ai_workbench_lib/modules/

file_discovery.rs

use anyhow::{Result};
use aws_sdk_s3::Client as S3Client;
use std::collections::HashMap;
use std::path::Path;
use super::file_splitter::types::FileType;
use tracing::{info, error, warn};

/// Handles discovery and processing of multiple files and folders
pub struct FileDiscovery {
    s3_client: S3Client,
    workspace_bucket: String,
}

impl FileDiscovery {
    pub fn new(s3_client: S3Client, workspace_bucket: String) -> Self {
        Self {
            s3_client,
            workspace_bucket,
        }
    }
    
    /// Discover all files to process from the input specification
    pub async fn discover_files(&self, input_spec: &str) -> Result<Vec<FileInfo>> {
        info!("=== FileDiscovery::discover_files DEBUG START ===");
        info!("Input spec: '{}'", input_spec);
        info!("Workspace bucket: '{}'", self.workspace_bucket);
        
        let result = if input_spec.ends_with('/') {
            info!("Detected folder input (ends with '/') - processing as folder");
            // Folder specified
            self.process_folder(input_spec).await
        } else {
            info!("Detected single file input (no trailing '/') - processing as single file");
            // Single file
            self.process_single_file(input_spec).await
        };
        
        match &result {
            Ok(files) => {
                info!("discover_files completed successfully with {} files", files.len());
            },
            Err(e) => {
                error!("discover_files failed: {}", e);
                error!("Error details: {:?}", e);
            }
        }
        
        info!("=== FileDiscovery::discover_files DEBUG END ===");
        result
    }
    
    /// Process a single file
    async fn process_single_file(&self, file_key: &str) -> Result<Vec<FileInfo>> {
        info!("=== process_single_file DEBUG START ===");
        info!("Processing single file: '{}'", file_key);
        info!("Full S3 path: s3://{}/{}", self.workspace_bucket, file_key);
        
        info!("Attempting to get file metadata...");
        let metadata_result = self.get_file_metadata(file_key).await;
        
        let metadata = match metadata_result {
            Ok(meta) => {
                info!("Metadata retrieved successfully:");
                info!("  - Size: {} bytes", meta.size_bytes);
                info!("  - File type: {:?}", meta.file_type);
                info!("  - Is processable: {}", meta.is_processable);
                
                if !meta.is_processable {
                    error!("File '{}' is not processable! Reasons:", file_key);
                    error!("  - File type: {:?}", meta.file_type);
                    error!("  - Size: {} bytes", meta.size_bytes);
                    if meta.size_bytes == 0 {
                        error!("  - File is empty (0 bytes)");
                    }
                    if matches!(meta.file_type, super::file_splitter::types::FileType::Binary) {
                        error!("  - File is binary type");
                        if meta.size_bytes > 100 * 1024 * 1024 {
                            error!("  - File is too large (> 100MB)");
                        }
                    }
                    return Err(anyhow::anyhow!("File '{}' is not processable", file_key));
                }
                
                meta
            },
            Err(e) => {
                error!("Failed to get metadata for file '{}': {}", file_key, e);
                error!("This usually means:");
                error!("  1. File does not exist at s3://{}/{}", self.workspace_bucket, file_key);
                error!("  2. No permissions to access the file");
                error!("  3. S3 service is unavailable");
                error!("  4. Incorrect bucket name or file path");
                return Err(e.context(format!("Failed to get metadata for single file: {}", file_key)));
            }
        };
        
        let file_info = FileInfo {
            s3_key: file_key.to_string(),
            relative_path: file_key.to_string(),
            size_bytes: metadata.size_bytes,
            file_type: metadata.file_type,
        };
        
        info!("Single file processed successfully:");
        info!("  - S3 key: '{}'", file_info.s3_key);
        info!("  - Relative path: '{}'", file_info.relative_path);
        info!("  - Size: {} bytes", file_info.size_bytes);
        info!("  - Type: {:?}", file_info.file_type);
        info!("=== process_single_file DEBUG END (SUCCESS) ===");
        
        Ok(vec![file_info])
    }
    
    
    /// Process a folder (recursively discover all files)
    async fn process_folder(&self, folder_prefix: &str) -> Result<Vec<FileInfo>> {
        info!("=== process_folder DEBUG START ===");
        info!("Processing folder: '{}'", folder_prefix);
        info!("Full S3 path: s3://{}/{}", self.workspace_bucket, folder_prefix);
        
        let mut files = Vec::new();
        let mut continuation_token: Option<String> = None;
        let mut request_count = 0;
        
        loop {
            request_count += 1;
            info!("Making S3 ListObjects request #{}", request_count);
            
            let mut request = self.s3_client
                .list_objects_v2()
                .bucket(&self.workspace_bucket)
                .prefix(folder_prefix)
                .max_keys(1000);
                
            if let Some(token) = &continuation_token {
                info!("Using continuation token: {}", token);
                request = request.continuation_token(token);
            }
            
            info!("Sending ListObjectsV2 request to S3...");
            let response_result = request.send().await;
            
            let response = match response_result {
                Ok(resp) => {
                    info!("S3 ListObjectsV2 request #{} successful!", request_count);
                    resp
                },
                Err(e) => {
                    error!("S3 ListObjectsV2 request failed: {}", e);
                    error!("Error details: {:?}", e);
                    error!("Possible causes:");
                    error!("  1. Folder/prefix '{}' does not exist in bucket '{}'", folder_prefix, self.workspace_bucket);
                    error!("  2. No list permissions for this bucket/prefix");
                    error!("  3. AWS credentials are invalid or expired");
                    error!("  4. Bucket '{}' does not exist", self.workspace_bucket);
                    error!("  5. Network connectivity issues");
                    return Err(anyhow::anyhow!("S3 ListObjectsV2 failed for folder '{}': {}", folder_prefix, e)
                        .context("Failed to list objects in S3"));
                }
            };
            
            let objects = response.contents();
            info!("Found {} objects in this batch", objects.len());
            
            if objects.is_empty() {
                warn!("No objects found in folder '{}' - this might be expected if folder is empty", folder_prefix);
            }
            
            for (i, object) in objects.iter().enumerate() {
                let key = object.key().unwrap_or_default();
                info!("Processing object {}: '{}'", i + 1, key);
                
                // Skip directories (keys ending with /)
                if key.ends_with('/') {
                    info!("  Skipping directory: '{}'", key);
                    continue;
                }
                
                info!("  Creating metadata for file: '{}'", key);
                let metadata = FileMetadata::from_s3_object(object, key);
                info!("  File metadata: size={} bytes, type={:?}, processable={}", 
                     metadata.size_bytes, metadata.file_type, metadata.is_processable);
                
                if metadata.is_processable {
                    let relative_path = key.strip_prefix(folder_prefix).unwrap_or(key).to_string();
                    let file_info = FileInfo {
                        s3_key: key.to_string(),
                        relative_path: relative_path.clone(),
                        size_bytes: metadata.size_bytes,
                        file_type: metadata.file_type,
                    };
                    
                    info!("  Adding processable file: '{}' (relative: '{}')", key, relative_path);
                    files.push(file_info);
                } else {
                    info!("  Skipping non-processable file: '{}' (size={}, type={:?})", 
                         key, metadata.size_bytes, metadata.file_type);
                }
            }
            
            continuation_token = response.next_continuation_token().map(|s| s.to_string());
            if let Some(ref token) = continuation_token {
                info!("More objects available, continuing with token: {}", token);
            } else {
                info!("No more objects to fetch, finishing folder processing");
                break;
            }
        }
        
        info!("Folder processing completed: {} processable files found", files.len());
        if files.is_empty() {
            warn!("No processable files found in folder '{}'", folder_prefix);
            warn!("This could mean:");
            warn!("  1. Folder is empty");
            warn!("  2. All files are binary/non-processable");
            warn!("  3. All files are empty (0 bytes)");
            warn!("  4. All files are too large (> 100MB for binary files)");
        }
        
        info!("=== process_folder DEBUG END ===");
        Ok(files)
    }
    
    /// Get metadata for a single file
    async fn get_file_metadata(&self, file_key: &str) -> Result<FileMetadata> {
        info!("=== get_file_metadata DEBUG START ===");
        info!("Getting metadata for file: '{}'", file_key);
        info!("Bucket: '{}'", self.workspace_bucket);
        info!("Full S3 path: s3://{}/{}", self.workspace_bucket, file_key);
        
        info!("Sending HEAD request to S3...");
        let response_result = self.s3_client
            .head_object()
            .bucket(&self.workspace_bucket)
            .key(file_key)
            .send()
            .await;
        
        let response = match response_result {
            Ok(resp) => {
                info!("S3 HEAD request successful!");
                info!("Response metadata:");
                if let Some(content_length) = resp.content_length() {
                    info!("  - Content length: {} bytes", content_length);
                } else {
                    warn!("  - No content length in response");
                }
                if let Some(content_type) = resp.content_type() {
                    info!("  - Content type: {}", content_type);
                } else {
                    info!("  - No content type specified");
                }
                if let Some(last_modified) = resp.last_modified() {
                    info!("  - Last modified: {:?}", last_modified);
                }
                resp
            },
            Err(e) => {
                error!("S3 HEAD request failed for '{}': {}", file_key, e);
                error!("Error details: {:?}", e);
                error!("Possible causes:");
                error!("  1. File '{}' does not exist in bucket '{}'", file_key, self.workspace_bucket);
                error!("  2. No read permissions for this file/bucket");
                error!("  3. AWS credentials are invalid or expired");
                error!("  4. Network connectivity issues");
                error!("  5. S3 service is temporarily unavailable");
                return Err(anyhow::anyhow!("S3 HEAD request failed for '{}': {}", file_key, e)
                    .context(format!("Failed to get metadata for {}", file_key)));
            }
        };
        
        let size_bytes = response.content_length().unwrap_or(0) as usize;
        let path = Path::new(file_key);
        let file_type = super::file_splitter::types::FileType::from_extension(path);
        let is_processable = Self::is_file_processable(&file_type, size_bytes);
        
        info!("File metadata analysis:");
        info!("  - Raw size from S3: {} bytes", size_bytes);
        info!("  - File path for type detection: {:?}", path);
        info!("  - Detected file type: {:?}", file_type);
        info!("  - Is processable: {}", is_processable);
        
        if !is_processable {
            info!("File is not processable because:");
            if size_bytes == 0 {
                info!("  - File is empty (0 bytes)");
            }
            if matches!(file_type, super::file_splitter::types::FileType::Binary) {
                info!("  - File type is Binary (not text-based)");
                if size_bytes > 100 * 1024 * 1024 {
                    info!("  - Binary file is too large (> 100MB): {} bytes", size_bytes);
                }
            }
        }
        
        let metadata = FileMetadata {
            size_bytes,
            file_type,
            is_processable,
        };
        
        info!("=== get_file_metadata DEBUG END (SUCCESS) ===");
        Ok(metadata)
    }
    
    /// Check if a file should be processed based on type and size
    fn is_file_processable(file_type: &FileType, size_bytes: usize) -> bool {
        
        // Skip binary files that are too large (> 100MB)
        if matches!(file_type, FileType::Binary) && size_bytes > 100 * 1024 * 1024 {
            return false;
        }
        
        // Skip empty files
        if size_bytes == 0 {
            return false;
        }
        
        // Process most text-based files
        !matches!(file_type, FileType::Binary)
    }
    
    /// Generate processing summary
    pub fn generate_processing_summary(&self, files: &[FileInfo]) -> ProcessingSummary {
        let mut summary = ProcessingSummary {
            total_files: files.len(),
            total_size_bytes: files.iter().map(|f| f.size_bytes).sum(),
            files_by_type: HashMap::new(),
            files_by_category: HashMap::new(),
        };
        
        for file in files {
            // Count by file type
            *summary.files_by_type.entry(file.file_type).or_insert(0) += 1;
            
            // Count by category
            let category = file.file_type.language_category();
            *summary.files_by_category.entry(category.to_string()).or_insert(0) += 1;
        }
        
        summary
    }
}

/// Information about a discovered file
#[derive(Debug, Clone,)]
pub struct FileInfo {
    pub s3_key: String,
    pub relative_path: String,
    pub size_bytes: usize,
    pub file_type: FileType,
}

/// File metadata
#[derive(Debug, Clone)]
struct FileMetadata {
    size_bytes: usize,
    pub file_type: FileType,
    is_processable: bool,
}

impl FileMetadata {
    fn from_s3_object(object: &aws_sdk_s3::types::Object, key: &str) -> Self {
        let size_bytes = object.size().unwrap_or(0) as usize;
        let path = Path::new(key);
        let file_type = FileType::from_extension(path);
        let is_processable = FileDiscovery::is_file_processable(&file_type, size_bytes);
        
        Self {
            size_bytes,
            file_type,
            is_processable,
        }
    }
}

/// Summary of files to be processed
#[derive(Debug, Clone)]
pub struct ProcessingSummary {
    pub total_files: usize,
    pub total_size_bytes: usize,
    pub files_by_type: HashMap<FileType, usize>,
    pub files_by_category: HashMap<String, usize>,
}

impl ProcessingSummary {
    pub fn format_summary(&self) -> String {
        let mut summary = format!(
            "Processing Summary:\n- Total files: {}\n- Total size: {:.2} MB\n",
            self.total_files,
            self.total_size_bytes as f64 / (1024.0 * 1024.0)
        );
        
        summary.push_str("\nFiles by category:\n");
        for (category, count) in &self.files_by_category {
            summary.push_str(&format!("- {}: {} files\n", category, count));
        }
        
        summary.push_str("\nFiles by type:\n");
        for (file_type, count) in &self.files_by_type {
            summary.push_str(&format!("- {:?}: {} files\n", file_type, count));
        }
        
        summary
    }
}


#[cfg(test)]
mod file_discovery_tests {
    use crate::modules::file_discovery::{FileDiscovery, FileInfo};
    use crate::modules::file_splitter::types::FileType;
    use aws_config::BehaviorVersion;
    use aws_sdk_s3::Client as S3Client;
    use std::env;

    /// Helper to create S3 client for testing
    async fn create_test_s3_client() -> S3Client {
        let config = aws_config::defaults(BehaviorVersion::latest())
            .region("eu-west-2")
            .load()
            .await;
        S3Client::new(&config)
    }

    /// Helper to get test workspace bucket from environment
    fn get_test_bucket() -> String {
        env::var("TEST_WORKSPACE_BUCKET")
            .unwrap_or_else(|_| "ai-workbench-6c9c43db-7fe6-42f1-8b11-8f82323f83f0-eu-west-2".to_string())
    }

    #[tokio::test]
    async fn test_single_file_discovery() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        // Test with a known file key (you'll provide this)
        let test_file_key = "test-data/sample.txt";
        
        let result = discovery.discover_files(test_file_key).await;
        
        match result {
            Ok(files) => {
                assert_eq!(files.len(), 1);
                let file = &files[0];
                assert_eq!(file.s3_key, test_file_key);
                assert_eq!(file.relative_path, test_file_key);
                assert!(file.size_bytes > 0);
                println!("Discovered file: {} ({} bytes, type: {:?})", 
                         file.s3_key, file.size_bytes, file.file_type);
            }
            Err(e) => {
                eprintln!("Test failed - file not found or error: {}", e);
                // For now, we'll make this non-failing until you provide real keys
                println!("Skipping test - no test file available");
            }
        }
    }

    #[tokio::test]
    async fn test_folder_discovery() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        // Test with a known folder prefix (you'll provide this)
        let test_folder = "text_files/";
        
        let result = discovery.discover_files(test_folder).await;
        
        match result {
            Ok(files) => {
                println!("Discovered {} files in folder '{}'", files.len(), test_folder);
                
                for file in &files {
                    assert!(file.s3_key.starts_with(test_folder));
                    assert!(!file.s3_key.ends_with('/'));  // Should not include directories
                    assert!(file.size_bytes > 0);
                    assert!(!file.relative_path.starts_with(test_folder));  // Should be relative
                    
                    println!("  - {} ({} bytes, type: {:?})", 
                             file.relative_path, file.size_bytes, file.file_type);
                }
                
                // Generate and display processing summary
                let summary = discovery.generate_processing_summary(&files);
                println!("\n{}", summary.format_summary());
            }
            Err(e) => {
                eprintln!("Test failed - folder not found or error: {}", e);
                println!("Skipping test - no test folder available");
            }
        }
    }

    #[tokio::test]
    async fn test_file_type_detection() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        // Test files with different extensions (you'll provide these)
        let test_files = vec![
            "test-data/sample.txt",
            "test-data/data.csv", 
            "test-data/config.json",
            "test-data/document.md",
            "test-data/script.py",
            "test-data/code.rs",
            "test-data/styles.css",
            "test-data/markup.html",
        ];

        for file_key in test_files {
            let result = discovery.discover_files(file_key).await;
            
            match result {
                Ok(files) if !files.is_empty() => {
                    let file = &files[0];
                    println!("File: {} -> Type: {:?}, Category: {}", 
                             file_key, file.file_type, file.file_type.language_category());
                    
                    // Verify file type matches extension
                    match file_key.split('.').last() {
                        Some("txt") | Some("md") => assert_eq!(file.file_type, FileType::Text),
                        Some("csv") => assert_eq!(file.file_type, FileType::Csv),
                        Some("json") => assert_eq!(file.file_type, FileType::Json),
                        Some("py") => assert_eq!(file.file_type, FileType::Python),
                        Some("rs") => assert_eq!(file.file_type, FileType::Rust),
                        Some("css") => assert_eq!(file.file_type, FileType::Css),
                        Some("html") | Some("htm") => assert_eq!(file.file_type, FileType::Html),
                        _ => {}  // Other types
                    }
                }
                Ok(_) => println!("File {} not found, skipping", file_key),
                Err(e) => println!("Error checking {}: {}", file_key, e),
            }
        }
    }

    #[tokio::test] 
    async fn test_large_folder_discovery() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        // Test with a folder that might have many files
        let large_folder = "large-dataset/";
        
        let result = discovery.discover_files(large_folder).await;
        
        match result {
            Ok(files) => {
                println!("Discovered {} files in large folder", files.len());
                
                if files.len() > 10 {
                    // Test pagination worked correctly
                    let total_size: usize = files.iter().map(|f| f.size_bytes).sum();
                    println!("Total size: {:.2} MB", total_size as f64 / (1024.0 * 1024.0));
                    
                    // Verify all files are unique
                    let mut keys: Vec<_> = files.iter().map(|f| &f.s3_key).collect();
                    keys.sort();
                    keys.dedup();
                    assert_eq!(keys.len(), files.len(), "Found duplicate files in discovery");
                }
                
                // Generate summary for large folder
                let summary = discovery.generate_processing_summary(&files);
                println!("\n{}", summary.format_summary());
            }
            Err(e) => {
                println!("Large folder test skipped: {}", e);
            }
        }
    }

    #[tokio::test]
    async fn test_file_filtering() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        // Test folder that might contain various file types including binaries
        let mixed_folder = "mixed-files/";
        
        let result = discovery.discover_files(mixed_folder).await;
        
        match result {
            Ok(files) => {
                println!("Found {} processable files in mixed folder", files.len());
                
                // All returned files should be processable
                for file in &files {
                    assert!(file.size_bytes > 0, "File {} has zero size", file.s3_key);
                    assert!(!matches!(file.file_type, FileType::Binary), 
                           "Binary file {} should have been filtered out", file.s3_key);
                    
                    // Very large files should be filtered (except text files)
                    if file.size_bytes > 100 * 1024 * 1024 {  // 100MB
                        assert!(!matches!(file.file_type, FileType::Binary),
                               "Large binary file {} should have been filtered", file.s3_key);
                    }
                }
            }
            Err(e) => {
                println!("Mixed files test skipped: {}", e);
            }
        }
    }

    #[tokio::test]
    async fn test_processing_summary() {
        let s3_client = create_test_s3_client().await;
        let bucket = get_test_bucket();
        let discovery = FileDiscovery::new(s3_client, bucket);

        let test_folder = "test-data/";
        
        let result = discovery.discover_files(test_folder).await;
        
        match result {
            Ok(files) if !files.is_empty() => {
                let summary = discovery.generate_processing_summary(&files);
                
                // Verify summary accuracy
                assert_eq!(summary.total_files, files.len());
                let expected_total_size: usize = files.iter().map(|f| f.size_bytes).sum();
                assert_eq!(summary.total_size_bytes, expected_total_size);
                
                // Verify file type counts
                let mut manual_type_counts = std::collections::HashMap::new();
                for file in &files {
                    *manual_type_counts.entry(file.file_type).or_insert(0) += 1;
                }
                assert_eq!(summary.files_by_type, manual_type_counts);
                
                // Verify category counts
                let mut manual_category_counts = std::collections::HashMap::new();
                for file in &files {
                    let category = file.file_type.language_category();
                    *manual_category_counts.entry(category.to_string()).or_insert(0) += 1;
                }
                assert_eq!(summary.files_by_category, manual_category_counts);
                
                // Test summary formatting
                let formatted = summary.format_summary();
                assert!(formatted.contains("Processing Summary:"));
                assert!(formatted.contains(&format!("Total files: {}", summary.total_files)));
                assert!(formatted.contains("Files by category:"));
                assert!(formatted.contains("Files by type:"));
                
                println!("Processing Summary Test Results:\n{}", formatted);
            }
            Ok(_) => println!("No files found for summary test"),
            Err(e) => println!("Summary test skipped: {}", e),
        }
    }

    #[test]
    fn test_file_info_structure() {
        use crate::modules::file_splitter::types::FileType;
        
        // Test FileInfo structure
        let file_info = FileInfo {
            s3_key: "test/path/file.txt".to_string(),
            relative_path: "path/file.txt".to_string(),
            size_bytes: 1024,
            file_type: FileType::Text,
        };
        
        assert_eq!(file_info.s3_key, "test/path/file.txt");
        assert_eq!(file_info.relative_path, "path/file.txt");
        assert_eq!(file_info.size_bytes, 1024);
        assert_eq!(file_info.file_type, FileType::Text);
        
        // Test that FileInfo can be cloned and debugged
        let cloned = file_info.clone();
        assert_eq!(file_info.s3_key, cloned.s3_key);
        
        let debug_str = format!("{:?}", file_info);
        assert!(debug_str.contains("FileInfo"));
    }
}