add Restore tool, fix metadata write error

src/app.rs

@@ -30,6 +30,24 @@ const FILE_NAME_LIST: &[u8] = include_bytes!("../assets/MHWs_STM_Release.list.zs
 const AUTO_CHUNK_SELECTION_SIZE_THRESHOLD: usize = 50 * 1024 * 1024; // 50MB
 const FALSE_TRUE_SELECTION: [&str; 2] = ["False", "True"];
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Mode {
+    Automatic = 0,
+    Manual = 1,
+    Restore = 2,
+}
+
+impl Mode {
+    fn from_index(index: usize) -> color_eyre::Result<Self> {
+        match index {
+            0 => Ok(Mode::Automatic),
+            1 => Ok(Mode::Manual),
+            2 => Ok(Mode::Restore),
+            _ => bail!("Invalid mode index: {index}"),
+        }
+    }
+}
+
 struct ChunkSelection {
     chunk_name: ChunkName,
     file_size: u64,
@@ -60,15 +78,15 @@ impl App {
         // Mode selection
         let mode = Select::with_theme(&ColorfulTheme::default())
             .with_prompt("Select mode")
-            .items(&["Automatic", "Manual"])
+            .items(&["Automatic", "Manual", "Restore"])
             .default(0)
             .interact()?;
-        let auto_mode = mode == 0;
+        let mode = Mode::from_index(mode)?;
 
-        if auto_mode {
+        match mode {
-            self.auto_mode()
+            Mode::Automatic => self.auto_mode(),
-        } else {
+            Mode::Manual => self.manual_mode(),
-            self.manual_mode()
+            Mode::Restore => self.restore_mode(),
         }
     }
 
@@ -111,13 +129,12 @@ impl App {
             .truncate(true)
             .write(true)
             .open(output_path)?;
-        let mut pak_writer = ree_pak_core::write::PakWriter::new(out_file, entries.len() as u64);
+        // +1 for metadata
+        let mut pak_writer =
+            ree_pak_core::write::PakWriter::new(out_file, (entries.len() as u64) + 1);
 
         // write metadata
-        let metadata = PakMetadata {
+        let metadata = PakMetadata::new(use_full_package_mode);
-            version: 1,
-            is_uncompressed_patch: true,
-        };
         metadata.write_to_pak(&mut pak_writer)?;
 
         let pak_writer_mtx = Arc::new(Mutex::new(pak_writer));
@@ -410,6 +427,174 @@ I'm sure I've checked the list, press Enter to continue"#,
 
         Ok(())
     }
+
+    fn restore_mode(&mut self) -> color_eyre::Result<()> {
+        let current_dir = std::env::current_dir()?;
+
+        let game_dir: String = Input::<String>::with_theme(&ColorfulTheme::default())
+            .show_default(true)
+            .default(current_dir.to_string_lossy().to_string())
+            .with_prompt("Input MonsterHunterWilds directory path")
+            .interact_text()
+            .unwrap()
+            .trim_matches(|c| c == '\"' || c == '\'')
+            .to_string();
+
+        let game_dir = Path::new(&game_dir);
+        if !game_dir.is_dir() {
+            bail!("game directory not exists.");
+        }
+
+        // scan all pak files, find files generated by this tool
+        println!("Scanning tool generated files...");
+        let dir = fs::read_dir(game_dir)?;
+        let mut tool_generated_files = Vec::new();
+        let mut backup_files = Vec::new();
+        let mut all_chunks = Vec::new();
+
+        for entry in dir {
+            let entry = entry?;
+            if !entry.file_type()?.is_file() {
+                continue;
+            }
+
+            let file_name = entry.file_name().to_string_lossy().to_string();
+            let file_path = entry.path();
+
+            // check backup files
+            if file_name.ends_with(".pak.backup") {
+                backup_files.push(file_path);
+                continue;
+            }
+
+            // check pak files
+            if !file_name.ends_with(".pak") || !file_name.starts_with("re_chunk_") {
+                continue;
+            }
+
+            // collect chunk info
+            if let Ok(chunk_name) = ChunkName::try_from_str(&file_name) {
+                all_chunks.push(chunk_name.clone());
+            }
+
+            // check if the file is generated by this tool
+            if let Ok(Some(metadata)) = self.check_tool_generated_file(&file_path) {
+                tool_generated_files.push((file_path, metadata));
+            }
+        }
+
+        if tool_generated_files.is_empty() && backup_files.is_empty() {
+            println!("No files found to restore.");
+            return Ok(());
+        }
+
+        println!(
+            "Found {} tool generated files and {} backup files",
+            tool_generated_files.len(),
+            backup_files.len()
+        );
+
+        // restore
+        let mut patch_files_to_remove = Vec::new();
+        for (file_path, metadata) in &tool_generated_files {
+            if metadata.is_full_package() {
+                // restore full package mode (replace mode)
+                // this is a replace mode generated file, find the corresponding backup file
+                let backup_path = file_path.with_extension("pak.backup");
+                if backup_path.exists() {
+                    println!("Restore replace mode file: {}", file_path.display());
+
+                    // delete the current file and restore the backup
+                    fs::remove_file(file_path)?;
+                    fs::rename(&backup_path, file_path)?;
+
+                    println!("   Restore backup file: {}", backup_path.display());
+                } else {
+                    println!("Warning: backup file not found {}", backup_path.display());
+                }
+            } else {
+                // restore patch mode
+                // this is a patch mode generated file
+                if let Ok(chunk_name) =
+                    ChunkName::try_from_str(&file_path.file_name().unwrap().to_string_lossy())
+                {
+                    patch_files_to_remove.push((file_path.clone(), chunk_name));
+                }
+            }
+        }
+
+        // remove patch files
+        if !patch_files_to_remove.is_empty() {
+            println!("Remove patch files...");
+
+            for (file_path, chunk_name) in patch_files_to_remove.iter().rev() {
+                println!("Remove patch file: {}", file_path.display());
+
+                // Check if there are any patches with higher numbers
+                let has_higher_patches = all_chunks.iter().any(|c| {
+                    c.major_id == chunk_name.major_id
+                        && c.sub_id == chunk_name.sub_id
+                        && match (c.sub_id, c.sub_patch_id) {
+                            (Some(_), Some(patch_id)) => {
+                                patch_id > chunk_name.sub_patch_id.unwrap()
+                            }
+                            (None, Some(patch_id)) => patch_id > chunk_name.patch_id.unwrap(),
+                            _ => false,
+                        }
+                });
+
+                if has_higher_patches {
+                    // create an empty patch file instead of deleting, to keep the patch sequence continuous
+                    self.create_empty_patch_file(file_path)?;
+                    println!("   Create empty patch file to keep sequence continuous");
+                } else {
+                    // no higher patches exist, safe to delete
+                    fs::remove_file(file_path)?;
+                    println!("   Removed patch file");
+                }
+            }
+        }
+
+        println!("Restore completed!");
+        Ok(())
+    }
+
+    /// check if the file is generated by this tool, return metadata
+    fn check_tool_generated_file(
+        &self,
+        file_path: &Path,
+    ) -> color_eyre::Result<Option<PakMetadata>> {
+        let file = match fs::File::open(file_path) {
+            Ok(file) => file,
+            Err(_) => return Ok(None),
+        };
+
+        let mut reader = io::BufReader::new(file);
+        let pak_archive = match ree_pak_core::read::read_archive(&mut reader) {
+            Ok(archive) => archive,
+            Err(_) => return Ok(None),
+        };
+
+        PakMetadata::from_pak_archive(reader, &pak_archive)
+    }
+
+    /// create an empty patch file
+    fn create_empty_patch_file(&self, file_path: &Path) -> color_eyre::Result<()> {
+        let out_file = OpenOptions::new()
+            .create(true)
+            .truncate(true)
+            .write(true)
+            .open(file_path)?;
+
+        let mut pak_writer = ree_pak_core::write::PakWriter::new(out_file, 1);
+
+        // write metadata to mark this is an empty patch file
+        let metadata = PakMetadata::new(false);
+        metadata.write_to_pak(&mut pak_writer)?;
+
+        pak_writer.finish()?;
+        Ok(())
+    }
 }
 
 fn is_tex_file(hash: u64, file_name_table: &FileNameTable) -> bool {

src/metadata.rs

@@ -14,11 +14,22 @@ const METADATA_KEY: &str = "__TEX_DECOMPRESSOR_METADATA__";
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PakMetadata {
-    pub version: u32,
+    version: u32,
-    pub is_uncompressed_patch: bool,
+    is_full_package: bool,
 }
 
 impl PakMetadata {
+    pub fn new(is_full_package: bool) -> Self {
+        Self {
+            version: 1,
+            is_full_package,
+        }
+    }
+
+    pub fn is_full_package(&self) -> bool {
+        self.is_full_package
+    }
+
     pub fn from_pak_archive<R>(
         reader: R,
         pak_archive: &PakArchive,