hikari-desktop/src-tauri/src/temp_manager.rs

use parking_lot::Mutex;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use uuid::Uuid;

const TEMP_DIR_NAME: &str = "hikari-uploads";

pub struct TempFileManager {
    base_dir: PathBuf,
    files: HashMap<String, Vec<PathBuf>>,
}

impl TempFileManager {
    pub fn new() -> Result<Self, String> {
        let base_dir = std::env::temp_dir().join(TEMP_DIR_NAME);

        if !base_dir.exists() {
            fs::create_dir_all(&base_dir)
                .map_err(|e| format!("Failed to create temp directory: {}", e))?;
        }

        Ok(TempFileManager {
            base_dir,
            files: HashMap::new(),
        })
    }

    #[allow(dead_code)]
    pub fn get_base_dir(&self) -> &Path {
        &self.base_dir
    }

    pub fn save_file(
        &mut self,
        conversation_id: &str,
        data: &[u8],
        original_filename: Option<&str>,
    ) -> Result<PathBuf, String> {
        let unique_id = Uuid::new_v4();
        let extension = original_filename
            .and_then(|name| Path::new(name).extension())
            .and_then(|ext| ext.to_str())
            .unwrap_or("bin");

        let filename = format!("{}_{}.{}", conversation_id, unique_id, extension);
        let file_path = self.base_dir.join(&filename);

        fs::write(&file_path, data)
            .map_err(|e| format!("Failed to write temp file: {}", e))?;

        self.files
            .entry(conversation_id.to_string())
            .or_default()
            .push(file_path.clone());

        Ok(file_path)
    }

    pub fn register_file(&mut self, conversation_id: &str, file_path: PathBuf) {
        self.files
            .entry(conversation_id.to_string())
            .or_default()
            .push(file_path);
    }

    pub fn get_files_for_conversation(&self, conversation_id: &str) -> Vec<PathBuf> {
        self.files
            .get(conversation_id)
            .cloned()
            .unwrap_or_default()
    }

    pub fn cleanup_conversation(&mut self, conversation_id: &str) -> Result<(), String> {
        if let Some(files) = self.files.remove(conversation_id) {
            for file_path in files {
                if file_path.exists() {
                    if let Err(e) = fs::remove_file(&file_path) {
                        eprintln!(
                            "Warning: Failed to remove temp file {:?}: {}",
                            file_path, e
                        );
                    }
                }
            }
        }
        Ok(())
    }

    pub fn cleanup_all(&mut self) -> Result<(), String> {
        let conversation_ids: Vec<String> = self.files.keys().cloned().collect();

        for conversation_id in conversation_ids {
            self.cleanup_conversation(&conversation_id)?;
        }

        Ok(())
    }

    pub fn cleanup_orphaned_files(&mut self) -> Result<usize, String> {
        let mut cleaned_count = 0;

        if !self.base_dir.exists() {
            return Ok(0);
        }

        let tracked_files: std::collections::HashSet<PathBuf> =
            self.files.values().flatten().cloned().collect();

        let entries = fs::read_dir(&self.base_dir)
            .map_err(|e| format!("Failed to read temp directory: {}", e))?;

        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_file() && !tracked_files.contains(&path) {
                if let Err(e) = fs::remove_file(&path) {
                    eprintln!("Warning: Failed to remove orphaned file {:?}: {}", path, e);
                } else {
                    cleaned_count += 1;
                }
            }
        }

        Ok(cleaned_count)
    }
}

impl Default for TempFileManager {
    fn default() -> Self {
        Self::new().expect("Failed to create TempFileManager")
    }
}

pub type SharedTempFileManager = Arc<Mutex<TempFileManager>>;

pub fn create_shared_temp_manager() -> Result<SharedTempFileManager, String> {
    Ok(Arc::new(Mutex::new(TempFileManager::new()?)))
}

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

    // Helper to create a TempFileManager with a custom base directory for testing
    fn create_test_manager(base_dir: PathBuf) -> TempFileManager {
        if !base_dir.exists() {
            fs::create_dir_all(&base_dir).expect("Failed to create test temp dir");
        }
        TempFileManager {
            base_dir,
            files: HashMap::new(),
        }
    }

    #[test]
    fn test_new_creates_base_directory() {
        let manager = TempFileManager::new().expect("Failed to create TempFileManager");
        assert!(manager.base_dir.exists());
    }

    #[test]
    fn test_get_base_dir_returns_correct_path() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let manager = create_test_manager(base_path.clone());

        assert_eq!(manager.get_base_dir(), base_path.as_path());
    }

    #[test]
    fn test_save_file_creates_file_with_content() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"Hello, world!";
        let result = manager.save_file("conv-1", data, Some("test.txt"));

        assert!(result.is_ok());
        let file_path = result.unwrap();
        assert!(file_path.exists());

        let content = fs::read(&file_path).expect("Failed to read file");
        assert_eq!(content, data);
    }

    #[test]
    fn test_save_file_uses_correct_extension() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"test data";
        let result = manager.save_file("conv-1", data, Some("document.pdf"));

        assert!(result.is_ok());
        let file_path = result.unwrap();
        assert_eq!(file_path.extension().unwrap(), "pdf");
    }

    #[test]
    fn test_save_file_uses_bin_extension_when_no_filename() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"binary data";
        let result = manager.save_file("conv-1", data, None);

        assert!(result.is_ok());
        let file_path = result.unwrap();
        assert_eq!(file_path.extension().unwrap(), "bin");
    }

    #[test]
    fn test_register_file_tracks_file_path() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let file_path = PathBuf::from("/some/path/file.txt");
        manager.register_file("conv-1", file_path.clone());

        let files = manager.get_files_for_conversation("conv-1");
        assert_eq!(files.len(), 1);
        assert_eq!(files[0], file_path);
    }

    #[test]
    fn test_get_files_for_conversation_returns_empty_for_unknown() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let manager = create_test_manager(base_path);

        let files = manager.get_files_for_conversation("unknown-conv");
        assert!(files.is_empty());
    }

    #[test]
    fn test_get_files_for_conversation_returns_all_files() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"test";
        manager.save_file("conv-1", data, Some("file1.txt")).unwrap();
        manager.save_file("conv-1", data, Some("file2.txt")).unwrap();
        manager.save_file("conv-2", data, Some("file3.txt")).unwrap();

        let files_conv1 = manager.get_files_for_conversation("conv-1");
        let files_conv2 = manager.get_files_for_conversation("conv-2");

        assert_eq!(files_conv1.len(), 2);
        assert_eq!(files_conv2.len(), 1);
    }

    #[test]
    fn test_cleanup_conversation_removes_files() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"test";
        let file_path = manager.save_file("conv-1", data, Some("test.txt")).unwrap();
        assert!(file_path.exists());

        let result = manager.cleanup_conversation("conv-1");
        assert!(result.is_ok());
        assert!(!file_path.exists());
        assert!(manager.get_files_for_conversation("conv-1").is_empty());
    }

    #[test]
    fn test_cleanup_conversation_handles_missing_files() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        // Register a file that doesn't exist
        manager.register_file("conv-1", PathBuf::from("/nonexistent/file.txt"));

        // Should not error, just skip missing files
        let result = manager.cleanup_conversation("conv-1");
        assert!(result.is_ok());
    }

    #[test]
    fn test_cleanup_conversation_for_unknown_returns_ok() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let result = manager.cleanup_conversation("unknown-conv");
        assert!(result.is_ok());
    }

    #[test]
    fn test_cleanup_all_removes_all_files() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"test";
        let file1 = manager.save_file("conv-1", data, Some("f1.txt")).unwrap();
        let file2 = manager.save_file("conv-2", data, Some("f2.txt")).unwrap();

        assert!(file1.exists());
        assert!(file2.exists());

        let result = manager.cleanup_all();
        assert!(result.is_ok());

        assert!(!file1.exists());
        assert!(!file2.exists());
        assert!(manager.files.is_empty());
    }

    #[test]
    fn test_cleanup_orphaned_files_removes_untracked() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path.clone());

        // Create a tracked file
        let data = b"tracked";
        let tracked_path = manager.save_file("conv-1", data, Some("tracked.txt")).unwrap();

        // Create an untracked (orphaned) file directly in the temp directory
        let orphan_path = base_path.join("orphan.txt");
        fs::write(&orphan_path, b"orphan").expect("Failed to create orphan file");

        assert!(tracked_path.exists());
        assert!(orphan_path.exists());

        let result = manager.cleanup_orphaned_files();
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 1); // One orphan removed

        assert!(tracked_path.exists()); // Tracked file still exists
        assert!(!orphan_path.exists()); // Orphan removed
    }

    #[test]
    fn test_cleanup_orphaned_returns_zero_when_none() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let data = b"test";
        manager.save_file("conv-1", data, Some("test.txt")).unwrap();

        let result = manager.cleanup_orphaned_files();
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 0);
    }

    #[test]
    fn test_cleanup_orphaned_returns_zero_when_dir_missing() {
        let mut manager = TempFileManager {
            base_dir: PathBuf::from("/nonexistent/dir"),
            files: HashMap::new(),
        };

        let result = manager.cleanup_orphaned_files();
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 0);
    }

    #[test]
    fn test_default_creates_manager() {
        // Default should work as long as we can create temp directories
        let manager = TempFileManager::default();
        assert!(manager.base_dir.exists());
    }

    #[test]
    fn test_create_shared_temp_manager() {
        let result = create_shared_temp_manager();
        assert!(result.is_ok());

        let shared = result.unwrap();
        let manager = shared.lock();
        assert!(manager.base_dir.exists());
    }

    #[test]
    fn test_multiple_files_same_conversation() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        // Save multiple files to same conversation
        for i in 0..5 {
            let data = format!("content {}", i);
            manager
                .save_file("conv-1", data.as_bytes(), Some(&format!("file{}.txt", i)))
                .unwrap();
        }

        let files = manager.get_files_for_conversation("conv-1");
        assert_eq!(files.len(), 5);

        // Each file should have unique content
        for (i, file_path) in files.iter().enumerate() {
            let content = fs::read_to_string(file_path).expect("Failed to read");
            assert_eq!(content, format!("content {}", i));
        }
    }

    #[test]
    fn test_file_paths_contain_conversation_id() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().join("hikari-test");
        let mut manager = create_test_manager(base_path);

        let file_path = manager
            .save_file("my-conversation-id", b"test", Some("test.txt"))
            .unwrap();

        let filename = file_path.file_name().unwrap().to_str().unwrap();
        assert!(filename.starts_with("my-conversation-id_"));
    }
}