code2prompt_core 4.2.0

//! This module contains the functions for traversing the directory and processing the files.
use crate::configuration::Code2PromptConfig;
use crate::file_processor;
use crate::filter::{build_globset, should_include_file};
use crate::sort::{FileSortMethod, sort_files, sort_tree};
use crate::tokenizer::count_tokens;
use crate::util::strip_utf8_bom;
use anyhow::Result;
use content_inspector::{ContentType, inspect};
use ignore::WalkBuilder;
use log::debug;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use termtree::Tree;

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct EntryMetadata {
    pub is_dir: bool,
    pub is_symlink: bool,
}

impl From<&std::fs::Metadata> for EntryMetadata {
    fn from(meta: &std::fs::Metadata) -> Self {
        Self {
            is_dir: meta.is_dir(),
            is_symlink: meta.is_symlink(),
        }
    }
}

/// Represents a file entry with all its metadata and content
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileEntry {
    pub path: String,
    pub extension: String,
    pub code: String,
    pub token_count: usize,
    pub metadata: EntryMetadata,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mod_time: Option<u64>,
}

/// Represents a file that needs to be processed
#[derive(Debug, Clone)]
struct FileToProcess {
    /// Absolute path to the file
    absolute_path: PathBuf,
    /// Relative path from the root
    relative_path: PathBuf,
    /// File metadata
    metadata: std::fs::Metadata,
}

/// Traverses the directory and returns the string representation of the tree and the vector of file entries.
///
/// This function uses the provided configuration to determine which files to include, how to format them,
/// and how to structure the directory tree.
///
/// # Arguments
///
/// * `config` - Configuration object containing path, include/exclude patterns, and other settings
/// * `selection_engine` - Optional SelectionEngine for advanced file selection with user actions
///
/// # Returns
///
/// * `Result<(String, Vec<FileEntry>)>` - A tuple containing the string representation of the directory
///   tree and a vector of file entries
pub fn traverse_directory(
    config: &Code2PromptConfig,
    selection_engine: Option<&mut crate::selection::SelectionEngine>,
) -> Result<(String, Vec<FileEntry>)> {
    // Phase 1: Discovery - Build tree and collect files to process
    let (tree, files_to_process) = discover_files(config, selection_engine)?;

    // Phase 2: Processing - Process files in parallel
    let mut files = process_files_parallel(files_to_process, config)?;

    // Phase 3: Assembly - Sort and return results
    assemble_results(tree, &mut files, config)
}

/// Phase 1: Discovery - Walk directories, build tree, and collect files that need processing
///
/// This phase is sequential because:
/// - Directory walking is already optimized
/// - Tree building needs sequential structure
/// - Selection engine has caching that would need synchronization
fn discover_files(
    config: &Code2PromptConfig,
    mut selection_engine: Option<&mut crate::selection::SelectionEngine>,
) -> Result<(Tree<String>, Vec<FileToProcess>)> {
    let canonical_root_path = config.path.canonicalize()?;
    let parent_directory = display_name(&canonical_root_path);

    let include_globset = build_globset(&config.include_patterns);
    let exclude_globset = build_globset(&config.exclude_patterns);

    // Build the Walker
    let walker = WalkBuilder::new(&canonical_root_path)
        .hidden(!config.hidden)
        .git_ignore(!config.no_ignore)
        .follow_links(config.follow_symlinks)
        .build()
        .filter_map(|entry| entry.ok());

    // Build the Tree
    let mut tree = Tree::new(parent_directory.to_owned());
    let mut files_to_process = Vec::new();

    for entry in walker {
        let path = entry.path();
        if let Ok(relative_path) = path.strip_prefix(&canonical_root_path) {
            // Use SelectionEngine if available, otherwise fall back to pattern matching
            let entry_match = if let Some(engine) = selection_engine.as_mut() {
                engine.is_selected(relative_path)
            } else {
                should_include_file(relative_path, &include_globset, &exclude_globset)
            };

            // Directory Tree
            let include_in_tree = config.full_directory_tree || entry_match;

            if include_in_tree {
                let mut current_tree = &mut tree;
                for component in relative_path.components() {
                    let component_str = component.as_os_str().to_string_lossy().to_string();
                    current_tree = if let Some(pos) = current_tree
                        .leaves
                        .iter_mut()
                        .position(|child| child.root == component_str)
                    {
                        &mut current_tree.leaves[pos]
                    } else {
                        let new_tree = Tree::new(component_str.clone());
                        current_tree.leaves.push(new_tree);
                        current_tree.leaves.last_mut().unwrap()
                    };
                }
            }

            // Collect files for processing
            if path.is_file()
                && entry_match
                && let Ok(metadata) = entry.metadata()
            {
                files_to_process.push(FileToProcess {
                    absolute_path: path.to_path_buf(),
                    relative_path: relative_path.to_path_buf(),
                    metadata,
                });
            }
        }
    }

    Ok((tree, files_to_process))
}

/// Phase 2: Processing - Process files in parallel using rayon
///
/// This phase processes files in parallel:
/// - Read file contents (I/O bound)
/// - Process file content (CPU/I/O bound)
/// - Tokenize if enabled (CPU bound)
/// - Build FileEntry structures
fn process_files_parallel(
    files_to_process: Vec<FileToProcess>,
    config: &Code2PromptConfig,
) -> Result<Vec<FileEntry>> {
    // Process files in parallel with rayon
    let files: Vec<Option<FileEntry>> = files_to_process
        .par_iter()
        .map(|file_info| process_single_file(file_info, config))
        .collect();

    // Filter out None values (files that failed to process or were empty)
    Ok(files.into_iter().flatten().collect())
}

/// Read file with single-pass binary detection
///
/// Reads file incrementally: first 8KB for binary detection, then remainder if text.
fn read_file_with_binary_check(path: &Path, file_size: u64) -> std::io::Result<Option<Vec<u8>>> {
    const SAMPLE_SIZE: usize = 8192;

    let mut file = fs::File::open(path)?;
    let mut buffer = Vec::with_capacity(file_size.min(1024 * 1024 * 10) as usize); // Cap at 10MB initial allocation

    // Read first chunk for binary detection
    let bytes_to_read = SAMPLE_SIZE.min(file_size as usize);
    let mut sample_buffer = vec![0u8; bytes_to_read];
    file.read_exact(&mut sample_buffer)?;

    // Check if binary
    if inspect(&sample_buffer) == ContentType::BINARY {
        return Ok(None); // Return None for binary files
    }

    // It's text! Add sample to buffer and read the rest
    buffer.extend_from_slice(&sample_buffer);

    // Read remaining bytes if file is larger than sample
    if file_size > SAMPLE_SIZE as u64 {
        file.read_to_end(&mut buffer)?;
    }

    Ok(Some(buffer))
}

/// Process a single file and return its FileEntry representation
fn process_single_file(file_info: &FileToProcess, config: &Code2PromptConfig) -> Option<FileEntry> {
    let path = &file_info.absolute_path;
    let relative_path = &file_info.relative_path;
    let metadata = &file_info.metadata;

    let code_bytes = match read_file_with_binary_check(path, metadata.len()) {
        Ok(Some(bytes)) => bytes,
        Ok(None) => {
            debug!("Skipped binary file: {}", path.display());
            return None;
        }
        Err(e) => {
            debug!("Failed to read file {}: {}", path.display(), e);
            return None;
        }
    };

    let clean_bytes = strip_utf8_bom(&code_bytes);

    // Get appropriate processor for file extension
    let extension = path.extension().and_then(|ext| ext.to_str()).unwrap_or("");
    let processor = file_processor::get_processor_for_extension(extension);

    // Process file content
    let code = match processor.process(clean_bytes, path) {
        Ok(processed) => processed,
        Err(e) => {
            log::warn!(
                "File processing failed for {}: {}. Using raw text fallback.",
                path.display(),
                e
            );
            String::from_utf8_lossy(clean_bytes).into_owned()
        }
    };

    // Wrap code block
    let code_block = wrap_code_block(&code, extension, config.line_numbers, config.no_codeblock);

    // Filter empty or invalid files
    if code.trim().is_empty() || code.contains(char::REPLACEMENT_CHARACTER) {
        debug!("Excluded file (empty or invalid UTF-8): {}", path.display());
        return None;
    }

    // Build filepath
    let file_path = if config.absolute_path {
        path.to_string_lossy().to_string()
    } else {
        relative_path.to_string_lossy().to_string()
    };

    // Always calculate token count in parallel (amortized by I/O wait time)
    // This enables zero-overhead token counting regardless of display preferences
    let token_count = count_tokens(&code, &config.encoding);

    // Get modification time if date sorting is requested
    let mod_time = if let Some(method) = config.sort_method {
        if method == FileSortMethod::DateAsc || method == FileSortMethod::DateDesc {
            metadata
                .modified()
                .ok()
                .and_then(|mtime| mtime.duration_since(std::time::SystemTime::UNIX_EPOCH).ok())
                .map(|d| d.as_secs())
        } else {
            None
        }
    } else {
        None
    };

    debug!(target: "included_files", "Included file: {}", file_path);

    Some(FileEntry {
        path: file_path,
        extension: extension.to_string(),
        code: code_block,
        token_count,
        metadata: EntryMetadata::from(metadata),
        mod_time,
    })
}

/// Phase 3: Assembly - Sort results and return
fn assemble_results(
    mut tree: Tree<String>,
    files: &mut [FileEntry],
    config: &Code2PromptConfig,
) -> Result<(String, Vec<FileEntry>)> {
    // Sort tree and files
    sort_tree(&mut tree, config.sort_method);
    sort_files(files, config.sort_method);

    Ok((tree.to_string(), files.to_owned()))
}

/// Returns the file name or the string representation of the path.
///
/// # Arguments
///
/// * `p` - The path to label.
///
/// # Returns
///
/// * `String` - The file name or string representation of the path.
pub fn display_name<P: AsRef<Path>>(p: P) -> String {
    let path = p.as_ref();
    // File name if available
    if let Some(name) = path.file_name() {
        return name.to_string_lossy().into_owned();
    }
    // Current directory name
    if let Ok(cwd) = std::env::current_dir()
        && let Some(name) = cwd.file_name()
    {
        return name.to_string_lossy().into_owned();
    }
    // Fallback
    ".".to_string()
}

/// Wraps the code block with a delimiter and adds line numbers if required.
///
/// # Arguments
///
/// * `code` - The code block to wrap.
/// * `extension` - The file extension of the code block.
/// * `line_numbers` - Whether to add line numbers to the code.
/// * `no_codeblock` - Whether to not wrap the code block with a delimiter.
///
/// # Returns
///
/// * `String` - The wrapped code block.
pub fn wrap_code_block(
    code: &str,
    extension: &str,
    line_numbers: bool,
    no_codeblock: bool,
) -> String {
    let delimiter = "`".repeat(3);
    let mut code_with_line_numbers = String::new();

    if line_numbers {
        for (line_number, line) in code.lines().enumerate() {
            code_with_line_numbers.push_str(&format!("{:4} | {}\n", line_number + 1, line));
        }
    } else {
        code_with_line_numbers = code.to_string();
    }

    if no_codeblock {
        code_with_line_numbers
    } else {
        format!(
            "{}{}\n{}\n{}",
            delimiter, extension, code_with_line_numbers, delimiter
        )
    }
}