ForgeLM/tauri-app/src-tauri/src/model_files.rs

use std::path::Path;

use tauri::AppHandle;
use tauri_plugin_dialog::DialogExt;

use crate::model_library::paths::resolve_models_dir;

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct ModelFile {
    pub name: String,
    pub path: String,
}

#[derive(serde::Serialize)]
pub struct PickDirResult {
    pub dir: String,
    pub models: Vec<ModelFile>,
}

fn model_file_from_path(path: &Path) -> ModelFile {
    ModelFile {
        name: path
            .file_name()
            .unwrap_or_default()
            .to_string_lossy()
            .to_string(),
        path: path.to_string_lossy().to_string(),
    }
}

fn scan_model_dir(path: &Path) -> Result<Vec<ModelFile>, String> {
    if !path.exists() {
        return Err(format!("Directory not found: {}", path.display()));
    }
    if !path.is_dir() {
        return Err(format!("Not a directory: {}", path.display()));
    }

    // Default scan is top-level only (no recursion into subdirectories)
    scan_model_dir_top_level(path)
}

/// Scans only the immediate children of `path` (depth=1, no recursion).
/// Returns only files with a `.gguf` extension (case-insensitive), sorted by name.
fn scan_model_dir_top_level(path: &Path) -> Result<Vec<ModelFile>, String> {
    if !path.exists() {
        return Err(format!("Directory not found: {}", path.display()));
    }
    if !path.is_dir() {
        return Err(format!("Not a directory: {}", path.display()));
    }

    let entries = std::fs::read_dir(path).map_err(|e| e.to_string())?;
    let mut models: Vec<ModelFile> = entries
        .flatten()
        .map(|e| e.path())
        .filter(|p| {
            p.is_file()
                && p.extension()
                    .map(|ext| ext.eq_ignore_ascii_case("gguf"))
                    .unwrap_or(false)
        })
        .map(|p| model_file_from_path(&p))
        .collect();

    models.sort_by(|a, b| a.name.cmp(&b.name));
    Ok(models)
}

/// Scans only `models/text-llm/` for `.gguf` files.
/// All other directories are ignored.
fn scan_model_dir_for_chat(dir: &Path, models: &mut Vec<ModelFile>) -> Result<(), String> {
    let text_llm_dir = dir.join("text-llm");

    if !text_llm_dir.exists() {
        return Ok(());
    }

    if !text_llm_dir.is_dir() {
        return Err(format!("Not a directory: {}", text_llm_dir.display()));
    }

    let entries = std::fs::read_dir(text_llm_dir).map_err(|e| e.to_string())?;

    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_file()
            && path
                .extension()
                .map(|ext| ext.eq_ignore_ascii_case("gguf"))
                .unwrap_or(false)
        {
            models.push(model_file_from_path(&path));
        }
    }

    Ok(())
}

/// Recursively scans all subdirectories for .gguf files.
/// Kept for future use when deep scanning may be needed.
#[allow(dead_code)]
fn scan_model_dir_recursive(dir: &Path, models: &mut Vec<ModelFile>) -> Result<(), String> {
    let entries = std::fs::read_dir(dir).map_err(|e| e.to_string())?;

    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            // Skip hidden directories (e.g. .git, .cache)
            if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                if name.starts_with('.') {
                    continue;
                }
            }
            scan_model_dir_recursive(&path, models)?;
        } else if path.is_file()
            && path
                .extension()
                .map(|ext| ext.eq_ignore_ascii_case("gguf"))
                .unwrap_or(false)
        {
            models.push(model_file_from_path(&path));
        }
    }

    Ok(())
}

fn import_model_file(src: &Path, dest_dir: &Path) -> Result<ModelFile, String> {
    if !src.exists() {
        return Err(format!("Source file not found: {}", src.display()));
    }
    if !src.is_file() {
        return Err("Source is not a file".into());
    }
    if !src
        .extension()
        .unwrap_or_default()
        .eq_ignore_ascii_case("gguf")
    {
        return Err("Only .gguf files are supported".into());
    }

    if !dest_dir.exists() {
        std::fs::create_dir_all(dest_dir)
            .map_err(|e| format!("Failed to create models directory: {e}"))?;
    }

    let file_name = src.file_name().unwrap_or_default();
    let dest_path = dest_dir.join(file_name);

    if src == dest_path {
        return Ok(model_file_from_path(&dest_path));
    }

    std::fs::copy(src, &dest_path).map_err(|e| format!("Failed to copy model: {e}"))?;
    Ok(model_file_from_path(&dest_path))
}

#[tauri::command]
pub fn scan_models(dir: String) -> Result<Vec<ModelFile>, String> {
    let path = if dir.trim().is_empty() {
        resolve_models_dir(None)
    } else {
        resolve_models_dir(Some(&dir))
    };
    
    if !path.exists() {
        return Err(format!("Directory not found: {}", path.display()));
    }
    if !path.is_dir() {
        return Err(format!("Not a directory: {}", path.display()));
    }

    let mut models = Vec::new();
    scan_model_dir_for_chat(&path, &mut models)?;
    models.sort_by(|a, b| a.name.cmp(&b.name));
    Ok(models)
}

#[tauri::command]
pub async fn pick_models_dir(app: AppHandle) -> Result<PickDirResult, String> {
    let folder = app
        .dialog()
        .file()
        .set_title("Select models folder")
        .blocking_pick_folder();

    match folder {
        Some(path) => {
            let dir = path.to_string();
            let models = scan_model_dir(Path::new(&dir))?;
            Ok(PickDirResult { dir, models })
        }
        None => Err("cancelled".into()),
    }
}

#[tauri::command]
pub async fn pick_model_files(app: AppHandle) -> Result<Vec<ModelFile>, String> {
    let files = app
        .dialog()
        .file()
        .set_title("Select model files")
        .add_filter("GGUF models", &["gguf"])
        .blocking_pick_files();

    match files {
        Some(paths) => Ok(paths.iter().map(|path| model_file_from_path(Path::new(&path.to_string()))).collect()),
        None => Err("cancelled".into()),
    }
}

#[tauri::command]
pub fn import_model(src_path: String, dest_dir: String) -> Result<ModelFile, String> {
    import_model_file(Path::new(&src_path), Path::new(&dest_dir))
}

#[cfg(test)]
mod tests {
    use super::{import_model_file, scan_model_dir, scan_model_dir_for_chat};
    use std::path::PathBuf;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn temp_path(name: &str) -> PathBuf {
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system time")
            .as_nanos();
        std::env::temp_dir().join(format!("forgelm-{name}-{unique}"))
    }

    #[test]
    fn scan_model_dir_filters_and_sorts_gguf_files() {
        let dir = temp_path("scan-models");
        std::fs::create_dir_all(&dir).expect("create temp dir");
        std::fs::write(dir.join("zeta.gguf"), b"demo").expect("write zeta");
        std::fs::write(dir.join("alpha.gguf"), b"demo").expect("write alpha");
        std::fs::write(dir.join("notes.txt"), b"ignore").expect("write non-model file");

        let models = scan_model_dir(&dir).expect("scan model dir");
        let names: Vec<_> = models.into_iter().map(|model| model.name).collect();

        assert_eq!(names, vec!["alpha.gguf", "zeta.gguf"]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn scan_model_dir_ignores_directories_named_like_models() {
        let dir = temp_path("scan-model-dirs");
        std::fs::create_dir_all(dir.join("folder.gguf")).expect("create fake model dir");
        std::fs::write(dir.join("real.gguf"), b"demo").expect("write model file");

        let models = scan_model_dir(&dir).expect("scan model dir");
        let names: Vec<_> = models.into_iter().map(|model| model.name).collect();

        assert_eq!(names, vec!["real.gguf"]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn scan_model_dir_for_chat_only_scans_text_llm_directory() {
        let dir = temp_path("scan-models-text-llm");
        let embeddings = dir.join("embeddings");
        let audio = dir.join("audio");
        let text_llm = dir.join("text-llm");
        std::fs::create_dir_all(&embeddings).expect("create embeddings dir");
        std::fs::create_dir_all(&audio).expect("create audio dir");
        std::fs::create_dir_all(&text_llm).expect("create text-llm dir");

        std::fs::write(dir.join("top-level.gguf"), b"demo").expect("write top-level model");
        std::fs::write(embeddings.join("embedding-model.gguf"), b"demo")
            .expect("write embedding model");
        std::fs::write(audio.join("whisper.gguf"), b"demo").expect("write audio model");
        std::fs::write(text_llm.join("chat-model.gguf"), b"demo").expect("write chat model");

        let mut models = Vec::new();
        scan_model_dir_for_chat(&dir, &mut models).expect("scan for chat");
        let names: Vec<_> = models.into_iter().map(|m| m.name).collect();

        assert_eq!(names, vec!["chat-model.gguf"]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn scan_model_dir_recursive_finds_nested_gguf() {
        // Tests that the recursive scan function still works for future use
        use super::scan_model_dir_recursive;
        
        let dir = temp_path("scan-models-recursive-fn");
        let sub = dir.join("mradermacher").join("Suri-Qwen-3.5-4B-Uncensored-i1-GGUF");
        std::fs::create_dir_all(&sub).expect("create nested dirs");
        std::fs::write(dir.join("top-level.gguf"), b"demo").expect("write top-level model");
        std::fs::write(sub.join("nested-model.gguf"), b"demo").expect("write nested model");

        let mut models = Vec::new();
        scan_model_dir_recursive(&dir, &mut models).expect("recursive scan");
        models.sort_by(|a, b| a.name.cmp(&b.name));

        let names: Vec<_> = models.iter().map(|m| m.name.as_str()).collect();
        assert_eq!(names, vec!["nested-model.gguf", "top-level.gguf"]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn scan_model_dir_skips_hidden_directories() {
        let dir = temp_path("scan-models-hidden");
        let hidden = dir.join(".git");
        std::fs::create_dir_all(&hidden).expect("create hidden dir");
        std::fs::write(dir.join("visible.gguf"), b"demo").expect("write visible model");
        std::fs::write(hidden.join("secret.gguf"), b"demo").expect("write hidden model");

        let models = scan_model_dir(&dir).expect("scan model dir");
        let names: Vec<_> = models.into_iter().map(|m| m.name).collect();

        assert_eq!(names, vec!["visible.gguf"]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn import_model_file_returns_existing_model_when_source_matches_destination() {
        let dir = temp_path("import-model");
        std::fs::create_dir_all(&dir).expect("create temp dir");
        let model_path = dir.join("demo.gguf");
        std::fs::write(&model_path, b"demo").expect("write demo model");

        let model = import_model_file(&model_path, &dir).expect("import model file");

        assert_eq!(model.name, "demo.gguf");
        assert_eq!(PathBuf::from(model.path), model_path);

        let _ = std::fs::remove_dir_all(&dir);
    }
}