diskforge_core/

file_finder.rs

use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};

use ignore::WalkBuilder;

/// Categories of file types for filtering.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FileCategory {
    DiskImage,
    Archive,
    Video,
    Installer,
    VmImage,
    Document,
    Other,
}

impl std::fmt::Display for FileCategory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::DiskImage => write!(f, "Disk Image"),
            Self::Archive => write!(f, "Archive"),
            Self::Video => write!(f, "Video"),
            Self::Installer => write!(f, "Installer"),
            Self::VmImage => write!(f, "VM Image"),
            Self::Document => write!(f, "Document"),
            Self::Other => write!(f, "Other"),
        }
    }
}

/// A file discovered by the finder.
#[derive(Debug, Clone)]
pub struct FoundFile {
    pub path: PathBuf,
    pub size: u64,
    pub mtime: Option<SystemTime>,
    pub category: FileCategory,
    pub is_downloaded: bool,
    pub download_source: Option<String>,
}

impl PartialEq for FoundFile {
    fn eq(&self, other: &Self) -> bool {
        self.size == other.size && self.path == other.path
    }
}

impl Eq for FoundFile {}

impl PartialOrd for FoundFile {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for FoundFile {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.size
            .cmp(&other.size)
            .then_with(|| self.path.cmp(&other.path))
    }
}

/// Options for the file finder.
#[derive(Debug, Clone)]
pub struct FindOptions {
    /// Root directories to scan.
    pub root_paths: Vec<PathBuf>,
    /// Minimum file size in bytes (files smaller than this are skipped).
    pub min_size: Option<u64>,
    /// Only include files older than this duration (based on mtime).
    pub older_than: Option<Duration>,
    /// Filter to specific file type categories.
    pub file_types: Option<Vec<FileCategory>>,
    /// Limit results to the top N largest files (uses BinaryHeap for memory efficiency).
    pub max_results: Option<usize>,
    /// Whether to validate file types with magic bytes (default: true).
    pub validate_types: bool,
    /// Additional paths to exclude (beyond the defaults).
    pub extra_exclusions: Vec<String>,
}

impl Default for FindOptions {
    fn default() -> Self {
        Self {
            root_paths: Vec::new(),
            min_size: None,
            older_than: None,
            file_types: None,
            max_results: None,
            validate_types: true,
            extra_exclusions: Vec::new(),
        }
    }
}

/// Paths that are always excluded from scanning.
const DEFAULT_EXCLUSIONS: &[&str] = &[
    "/System",
    "/usr/bin",
    "/usr/lib",
    "/usr/libexec",
    "/usr/sbin",
    "/usr/share",
    "/bin",
    "/sbin",
    "/private/var/db",
];

/// Components that signal we are inside an excluded tree.
const EXCLUDED_COMPONENTS: &[&str] = &[".git"];

/// Check if a path is inside a .app bundle (should be treated as atomic).
fn is_inside_app_bundle(path: &Path) -> bool {
    for ancestor in path.ancestors().skip(1) {
        if ancestor.extension().is_some_and(|e| e == "app") {
            return true;
        }
    }
    false
}

/// Check if a path should be excluded from scanning.
fn should_exclude(path: &Path, extra_exclusions: &[String]) -> bool {
    let path_str = path.to_string_lossy();

    // Check default absolute exclusions
    for excl in DEFAULT_EXCLUSIONS {
        if path_str.starts_with(excl) {
            return true;
        }
    }

    // Check extra exclusions
    for excl in extra_exclusions {
        if path_str.contains(excl.as_str()) {
            return true;
        }
    }

    // Check component-based exclusions
    for component in path.components() {
        let s = component.as_os_str().to_string_lossy();
        for excl in EXCLUDED_COMPONENTS {
            if s == *excl {
                return true;
            }
        }
    }

    // Check if inside .app bundle
    if is_inside_app_bundle(path) {
        return true;
    }

    false
}

/// Detect file category by extension.
pub fn detect_by_extension(path: &Path) -> FileCategory {
    let name = path
        .file_name()
        .unwrap_or(OsStr::new(""))
        .to_string_lossy()
        .to_lowercase();

    // Check compound extensions first
    if name.ends_with(".tar.gz") || name.ends_with(".tar.bz2") || name.ends_with(".tar.xz") {
        return FileCategory::Archive;
    }

    let ext = path
        .extension()
        .unwrap_or(OsStr::new(""))
        .to_string_lossy()
        .to_lowercase();

    match ext.as_str() {
        // Disk images
        "dmg" | "iso" | "img" | "sparseimage" | "sparsebundle" => FileCategory::DiskImage,
        // Archives
        "zip" | "tar" | "tgz" | "rar" | "7z" | "xz" | "gz" | "bz2" => FileCategory::Archive,
        // Videos
        "mp4" | "mov" | "avi" | "mkv" | "wmv" | "flv" | "m4v" | "webm" => FileCategory::Video,
        // Installers
        "pkg" | "mpkg" => FileCategory::Installer,
        // VM images
        "vmdk" | "vdi" | "qcow2" | "vhd" => FileCategory::VmImage,
        // Documents (large ones)
        "pdf" => FileCategory::Document,
        _ => FileCategory::Other,
    }
}

/// Validate file type using magic bytes via the `infer` crate.
/// If magic bytes disagree with extension-based detection, trust magic bytes.
/// If infer returns None (unknown), keep the extension-based category.
///
/// Exception: DMG files have the `koly` signature at EOF-512, not at the header,
/// so `infer` misclassifies them as Archive. For `.dmg` files, trust the extension.
pub fn validate_by_magic_bytes(path: &Path, extension_category: FileCategory) -> FileCategory {
    // DMG files: koly magic is at EOF-512, infer reads header and misclassifies as Archive.
    // Trust the extension-based detection for .dmg files.
    if extension_category == FileCategory::DiskImage {
        let ext = path
            .extension()
            .unwrap_or_default()
            .to_string_lossy()
            .to_lowercase();
        if ext == "dmg" {
            return FileCategory::DiskImage;
        }
    }

    let Ok(kind) = infer::get_from_path(path) else {
        return extension_category;
    };
    let Some(kind) = kind else {
        return extension_category;
    };

    // Map infer MIME types to our categories
    let mime = kind.mime_type();
    if mime.starts_with("video/") {
        return FileCategory::Video;
    }
    match mime {
        "application/x-apple-diskimage" => FileCategory::DiskImage,
        "application/zip"
        | "application/gzip"
        | "application/x-tar"
        | "application/x-rar-compressed"
        | "application/x-7z-compressed"
        | "application/x-bzip2"
        | "application/x-xz"
        | "application/zstd" => FileCategory::Archive,
        "application/pdf" => FileCategory::Document,
        "application/x-xar" => FileCategory::Installer, // .pkg files use xar format
        _ => extension_category,
    }
}

/// Read the quarantine xattr to detect downloaded files.
/// Returns (is_downloaded, download_source).
pub fn check_quarantine(path: &Path) -> (bool, Option<String>) {
    match xattr::get(path, "com.apple.quarantine") {
        Ok(Some(value)) => {
            let value_str = String::from_utf8_lossy(&value);
            // Format: flag;timestamp;agent_name;UUID
            let parts: Vec<&str> = value_str.split(';').collect();
            let source = parts
                .get(2)
                .map(|s| s.to_string())
                .filter(|s| !s.is_empty());
            (true, source)
        }
        _ => (false, None),
    }
}

/// Shared counter for progress reporting during parallel walk.
pub struct ScanProgress {
    pub files_scanned: std::sync::atomic::AtomicU64,
}

impl Default for ScanProgress {
    fn default() -> Self {
        Self::new()
    }
}

impl ScanProgress {
    pub fn new() -> Self {
        Self {
            files_scanned: std::sync::atomic::AtomicU64::new(0),
        }
    }

    pub fn count(&self) -> u64 {
        self.files_scanned
            .load(std::sync::atomic::Ordering::Relaxed)
    }
}

/// Find files matching the given options using parallel traversal.
///
/// Uses `ignore::WalkParallel` for 3-6x speedup over single-threaded traversal.
/// Files are filtered by size during the walk to avoid collecting small files.
/// Age is determined by `mtime` (NOT atime -- APFS disables atime by default).
///
/// Returns results sorted by size descending.
pub fn find_files(options: &FindOptions, progress: Option<&ScanProgress>) -> Vec<FoundFile> {
    let now = SystemTime::now();
    let min_size = options.min_size.unwrap_or(0);
    let extra_exclusions = options.extra_exclusions.clone();

    // Build results collector based on whether we have a max_results limit
    let collector: Arc<Mutex<TopNCollector>> =
        Arc::new(Mutex::new(TopNCollector::new(options.max_results)));

    for root in &options.root_paths {
        if !root.exists() {
            continue;
        }

        let mut builder = WalkBuilder::new(root);
        builder
            .hidden(false) // Don't skip hidden files -- users may want to find them
            .git_ignore(true) // Respect .gitignore by default
            .git_global(false)
            .git_exclude(false)
            .follow_links(false) // Don't follow symlinks (avoid infinite loops)
            .threads(num_cpus());

        let walker = builder.build_parallel();
        let collector_ref = Arc::clone(&collector);
        let extra_excl = extra_exclusions.clone();
        let older_than = options.older_than;
        let file_types = options.file_types.clone();

        walker.run(|| {
            let collector = Arc::clone(&collector_ref);
            let extra_excl = extra_excl.clone();
            let file_types = file_types.clone();

            Box::new(move |entry| {
                let Ok(entry) = entry else {
                    return ignore::WalkState::Continue;
                };

                let path = entry.path();

                // Skip excluded paths early (prevents descent into excluded trees)
                if should_exclude(path, &extra_excl) {
                    return ignore::WalkState::Skip;
                }

                // Only process files (not directories)
                let Some(file_type) = entry.file_type() else {
                    return ignore::WalkState::Continue;
                };
                if !file_type.is_file() {
                    return ignore::WalkState::Continue;
                }

                // Increment progress counter
                if let Some(prog) = progress {
                    prog.files_scanned
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                }

                // Get metadata for size and mtime
                let Ok(meta) = entry.metadata() else {
                    return ignore::WalkState::Continue;
                };

                let size = meta.len();

                // Filter by minimum size
                if size < min_size {
                    return ignore::WalkState::Continue;
                }

                // Filter by age (mtime)
                if let Some(max_age) = older_than
                    && let Ok(mtime) = meta.modified()
                    && let Ok(age) = now.duration_since(mtime)
                    && age < max_age
                {
                    return ignore::WalkState::Continue;
                }

                // Detect category by extension (fast, no I/O)
                let category = detect_by_extension(path);

                // Filter by file type if specified
                if let Some(ref types) = file_types {
                    if !types.contains(&category) && category != FileCategory::Other {
                        return ignore::WalkState::Continue;
                    }
                    // If category is Other and we have type filters, still skip
                    if category == FileCategory::Other && !types.contains(&FileCategory::Other) {
                        return ignore::WalkState::Continue;
                    }
                }

                let mtime = meta.modified().ok();

                let file = FoundFile {
                    path: path.to_path_buf(),
                    size,
                    mtime,
                    category,
                    is_downloaded: false,
                    download_source: None,
                };

                if let Ok(mut coll) = collector.lock() {
                    coll.push(file);
                }

                ignore::WalkState::Continue
            })
        });
    }

    let inner = Arc::try_unwrap(collector)
        .map(|mutex| mutex.into_inner().unwrap_or_else(|e| e.into_inner()))
        .unwrap_or_else(|arc| {
            let lock = arc.lock().unwrap();
            lock.clone()
        });
    let mut results = inner.into_vec();

    // Second pass: validate types with magic bytes and check quarantine xattr
    if options.validate_types {
        for file in &mut results {
            file.category = validate_by_magic_bytes(&file.path, file.category);
        }
    }

    // Check quarantine xattr for download detection
    for file in &mut results {
        let (is_downloaded, source) = check_quarantine(&file.path);
        file.is_downloaded = is_downloaded;
        file.download_source = source;
    }

    // Re-filter by file type after magic byte validation (category may have changed)
    if let Some(ref types) = options.file_types {
        results.retain(|f| types.contains(&f.category));
    }

    // Sort by size descending
    results.sort_by(|a, b| b.size.cmp(&a.size));
    results
}

/// Memory-efficient top-N collector using a BinaryHeap (min-heap).
/// When max is set, keeps only the N largest items. O(n log k) instead of O(n log n).
/// When max is None, collects all items into a Vec.
#[derive(Clone)]
struct TopNCollector {
    max: Option<usize>,
    heap: BinaryHeap<Reverse<FoundFile>>,
    vec: Vec<FoundFile>,
}

impl TopNCollector {
    fn new(max: Option<usize>) -> Self {
        Self {
            max,
            heap: BinaryHeap::new(),
            vec: Vec::new(),
        }
    }

    fn push(&mut self, file: FoundFile) {
        match self.max {
            Some(n) if n > 0 => {
                self.heap.push(Reverse(file));
                if self.heap.len() > n {
                    self.heap.pop(); // Remove smallest
                }
            }
            _ => {
                self.vec.push(file);
            }
        }
    }

    fn into_vec(self) -> Vec<FoundFile> {
        match self.max {
            Some(_) => self
                .heap
                .into_sorted_vec()
                .into_iter()
                .map(|Reverse(f)| f)
                .collect(),
            None => self.vec,
        }
    }
}

/// Get a reasonable thread count for parallel walking.
fn num_cpus() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    fn setup_test_dir(name: &str) -> PathBuf {
        let tmp = std::env::temp_dir().join(format!("diskforge_test_ff_{name}"));
        let _ = fs::remove_dir_all(&tmp);
        fs::create_dir_all(&tmp).unwrap();
        tmp
    }

    fn create_file(dir: &Path, name: &str, size: usize) -> PathBuf {
        let path = dir.join(name);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        fs::write(&path, vec![0u8; size]).unwrap();
        path
    }

    #[test]
    fn detect_extension_disk_image() {
        assert_eq!(
            detect_by_extension(Path::new("file.dmg")),
            FileCategory::DiskImage
        );
        assert_eq!(
            detect_by_extension(Path::new("file.iso")),
            FileCategory::DiskImage
        );
        assert_eq!(
            detect_by_extension(Path::new("file.IMG")),
            FileCategory::DiskImage
        );
    }

    #[test]
    fn detect_extension_archive() {
        assert_eq!(
            detect_by_extension(Path::new("file.zip")),
            FileCategory::Archive
        );
        assert_eq!(
            detect_by_extension(Path::new("file.tar.gz")),
            FileCategory::Archive
        );
        assert_eq!(
            detect_by_extension(Path::new("file.7z")),
            FileCategory::Archive
        );
        assert_eq!(
            detect_by_extension(Path::new("file.rar")),
            FileCategory::Archive
        );
    }

    #[test]
    fn detect_extension_video() {
        assert_eq!(
            detect_by_extension(Path::new("file.mp4")),
            FileCategory::Video
        );
        assert_eq!(
            detect_by_extension(Path::new("file.mkv")),
            FileCategory::Video
        );
        assert_eq!(
            detect_by_extension(Path::new("movie.MOV")),
            FileCategory::Video
        );
    }

    #[test]
    fn detect_extension_installer() {
        assert_eq!(
            detect_by_extension(Path::new("setup.pkg")),
            FileCategory::Installer
        );
        assert_eq!(
            detect_by_extension(Path::new("setup.mpkg")),
            FileCategory::Installer
        );
    }

    #[test]
    fn detect_extension_vm_image() {
        assert_eq!(
            detect_by_extension(Path::new("disk.vmdk")),
            FileCategory::VmImage
        );
        assert_eq!(
            detect_by_extension(Path::new("disk.qcow2")),
            FileCategory::VmImage
        );
    }

    #[test]
    fn detect_extension_document() {
        assert_eq!(
            detect_by_extension(Path::new("doc.pdf")),
            FileCategory::Document
        );
    }

    #[test]
    fn detect_extension_other() {
        assert_eq!(
            detect_by_extension(Path::new("file.txt")),
            FileCategory::Other
        );
        assert_eq!(
            detect_by_extension(Path::new("file.rs")),
            FileCategory::Other
        );
        assert_eq!(detect_by_extension(Path::new("noext")), FileCategory::Other);
    }

    #[test]
    fn find_files_size_filter() {
        let tmp = setup_test_dir("size_filter");
        create_file(&tmp, "small.zip", 100);
        create_file(&tmp, "big.zip", 10_000);
        create_file(&tmp, "huge.dmg", 100_000);

        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(5_000),
            validate_types: false,
            ..Default::default()
        };

        let results = find_files(&options, None);
        assert_eq!(results.len(), 2, "Should find 2 files above 5KB");
        assert!(
            results[0].size >= results[1].size,
            "Should be sorted by size desc"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn find_files_max_results() {
        let tmp = setup_test_dir("max_results");
        for i in 0..10 {
            create_file(&tmp, &format!("file{i}.zip"), (i + 1) * 1000);
        }

        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            max_results: Some(3),
            validate_types: false,
            ..Default::default()
        };

        let results = find_files(&options, None);
        assert_eq!(results.len(), 3, "Should return exactly 3 results");
        // Should be the 3 largest
        assert!(results[0].size >= results[1].size);
        assert!(results[1].size >= results[2].size);

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn find_files_type_filter() {
        let tmp = setup_test_dir("type_filter");
        create_file(&tmp, "movie.mp4", 5000);
        create_file(&tmp, "archive.zip", 5000);
        create_file(&tmp, "image.dmg", 5000);

        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            file_types: Some(vec![FileCategory::Video]),
            validate_types: false,
            ..Default::default()
        };

        let results = find_files(&options, None);
        assert_eq!(results.len(), 1, "Should find only the video");
        assert_eq!(results[0].category, FileCategory::Video);

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn find_files_excludes_git_objects() {
        let tmp = setup_test_dir("git_objects");
        create_file(&tmp, ".git/objects/pack/bigpack.zip", 10_000);
        create_file(&tmp, "normal.zip", 10_000);

        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            validate_types: false,
            ..Default::default()
        };

        let results = find_files(&options, None);
        assert_eq!(results.len(), 1, "Should exclude .git/objects file");
        assert!(!results[0].path.to_string_lossy().contains(".git"));

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn find_files_excludes_app_bundle_contents() {
        let tmp = setup_test_dir("app_bundle");
        create_file(&tmp, "SomeApp.app/Contents/MacOS/binary", 10_000);
        create_file(&tmp, "outside.dmg", 10_000);

        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            validate_types: false,
            ..Default::default()
        };

        let results = find_files(&options, None);
        assert_eq!(results.len(), 1, "Should exclude file inside .app bundle");
        assert!(results[0].path.to_string_lossy().contains("outside.dmg"));

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn find_files_progress_counter() {
        let tmp = setup_test_dir("progress");
        create_file(&tmp, "a.zip", 1000);
        create_file(&tmp, "b.zip", 2000);
        create_file(&tmp, "c.zip", 3000);

        let progress = ScanProgress::new();
        let options = FindOptions {
            root_paths: vec![tmp.clone()],
            min_size: Some(1),
            validate_types: false,
            ..Default::default()
        };

        find_files(&options, Some(&progress));
        assert!(
            progress.count() >= 3,
            "Should have scanned at least 3 files"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn dmg_not_misclassified_as_archive() {
        // Create a fake .dmg file (infer would see arbitrary bytes, not the koly trailer)
        let tmp = setup_test_dir("dmg_fix");
        let dmg_path = create_file(&tmp, "test.dmg", 5000);

        // Extension-based should be DiskImage
        assert_eq!(
            detect_by_extension(&dmg_path),
            FileCategory::DiskImage,
            "Extension detection should return DiskImage for .dmg"
        );

        // Magic-byte validation should NOT override .dmg to Archive
        let validated = validate_by_magic_bytes(&dmg_path, FileCategory::DiskImage);
        assert_eq!(
            validated,
            FileCategory::DiskImage,
            "DMG files must stay as DiskImage even after magic-byte validation"
        );

        fs::remove_dir_all(&tmp).ok();
    }

    #[test]
    fn top_n_collector_unlimited() {
        let mut coll = TopNCollector::new(None);
        for i in 0..5 {
            coll.push(FoundFile {
                path: PathBuf::from(format!("file{i}")),
                size: (i + 1) as u64 * 100,
                mtime: None,
                category: FileCategory::Other,
                is_downloaded: false,
                download_source: None,
            });
        }
        let results = coll.into_vec();
        assert_eq!(results.len(), 5);
    }

    #[test]
    fn top_n_collector_limited() {
        let mut coll = TopNCollector::new(Some(2));
        for i in 0..5 {
            coll.push(FoundFile {
                path: PathBuf::from(format!("file{i}")),
                size: (i + 1) as u64 * 100,
                mtime: None,
                category: FileCategory::Other,
                is_downloaded: false,
                download_source: None,
            });
        }
        let results = coll.into_vec();
        assert_eq!(results.len(), 2);
        // Should be the 2 largest (400 and 500)
        assert!(results.iter().all(|f| f.size >= 400));
    }
}