ra2-mix 0.0.2

//! Writer module for RA2 MIX files

use super::*;
use std::{collections::VecDeque, fs::create_dir_all, path::PathBuf};

impl MixPackage {
    /// # Arguments
    ///
    /// * `output`:
    ///
    /// # Examples
    ///
    /// ```
    /// ```
    pub fn save(self, output: &Path) -> Result<usize, Ra2Error> {
        if !output.is_file() {
            return Err(Ra2Error::FileNotFound("must be a file".to_string()));
        }
        let data = self.encode()?;
        std::fs::write(output, &data)?;
        Ok(data.len())
    }
    pub fn dump(self, output: &Path) -> Result<usize, Ra2Error> {
        if !output.exists() {
            println!("Skipping file {}", output.display());
            create_dir_all(output)?;
        }
        if !output.is_dir() {
            return Err(Ra2Error::FileNotFound("file exists but not folder".to_string()));
        }
        unsafe { self.dump_unchecked(output) }
    }
    pub unsafe fn dump_unchecked(self, output: &Path) -> Result<usize, Ra2Error> {
        for (filename, data) in self.files.iter() {
            let mut file = File::create(output.join(&filename))?;
            file.write_all(&data)?;
        }
        Ok(self.files.len())
    }
    /// # Arguments
    ///
    /// * `output`:
    ///
    /// # Examples
    ///
    /// ```
    /// ```
    pub fn encode(self) -> Result<Vec<u8>, Ra2Error> {
        let file_map = coalesce_input_files(self.game, &self.files)?;

        // Create file information list
        let mut file_information_list: Vec<FileInfo> =
            file_map.iter().map(|(filename, data)| FileInfo { file_id: ra2_crc(filename), data: data.clone() }).collect();

        // Sort by file ID
        file_information_list.sort_by_key(|file_info| file_info.file_id);

        // Generate file entries and body data
        let mut offset = 0u32;
        let mut file_entry_data = Vec::new();
        let mut body_data = Vec::new();

        for file_info in &file_information_list {
            let size = file_info.data.len() as u32;

            // Write file entry
            file_entry_data.write_u32::<LittleEndian>(file_info.file_id)?;
            file_entry_data.write_u32::<LittleEndian>(offset)?;
            file_entry_data.write_u32::<LittleEndian>(size)?;

            // Write file data
            body_data.extend_from_slice(&file_info.data);

            offset += size;
        }

        // Combine all parts
        let mut mix_data = create_mix_header(&file_map)?;
        mix_data.extend_from_slice(&file_entry_data);
        mix_data.extend_from_slice(&body_data);

        Ok(mix_data)
    }
}

/// Creates MIX database data
fn get_mix_db_data(filenames: &[String], game: CncGame) -> Vec<u8> {
    let num_files = filenames.len();
    let db_size_in_bytes = XCC_HEADER_SIZE + filenames.iter().map(|filename| filename.len() + 1).sum::<usize>();

    let mut bytes_data = Vec::with_capacity(db_size_in_bytes);

    // Write XCC ID bytes
    bytes_data.extend_from_slice(XCC_ID_BYTES);
    // Pad to 32 bytes
    bytes_data.resize(32, 0);

    // Write header fields
    bytes_data.write_u32::<LittleEndian>(db_size_in_bytes as u32).unwrap();
    bytes_data.write_u32::<LittleEndian>(XCC_FILE_TYPE).unwrap();
    bytes_data.write_u32::<LittleEndian>(XCC_FILE_VERSION).unwrap();
    bytes_data.write_u32::<LittleEndian>(game as u32).unwrap();
    bytes_data.write_u32::<LittleEndian>(num_files as u32).unwrap();

    // Write filenames with null terminators
    for filename in filenames {
        bytes_data.extend_from_slice(filename.as_bytes());
        bytes_data.push(0); // Null terminator
    }

    bytes_data
}

/// Processes input files and creates a file map
pub fn coalesce_input_files(game: CncGame, file_map: &HashMap<String, Vec<u8>>) -> Result<HashMap<String, Vec<u8>>, Ra2Error> {
    let mut extra_file_map = file_map.clone();
    // Get filenames and create mix database
    let mut filenames: Vec<String> = extra_file_map.keys().cloned().collect();
    filenames.push(MIX_DB_FILENAME.to_string());
    // Sort filenames
    let db_data = get_mix_db_data(&filenames, game);
    extra_file_map.insert(MIX_DB_FILENAME.to_string(), db_data);
    Ok(extra_file_map)
}

/// Creates a MIX file header
fn create_mix_header(file_map: &HashMap<String, Vec<u8>>) -> Result<Vec<u8>, Ra2Error> {
    let flags = 0u32;
    let file_count = file_map.len() as u16;
    let data_size = file_map.values().map(|data| data.len() as u32).sum();

    let mut header = Vec::with_capacity(HEADER_SIZE);
    header.write_u32::<LittleEndian>(flags)?;
    header.write_u16::<LittleEndian>(file_count)?;
    header.write_u32::<LittleEndian>(data_size)?;

    Ok(header)
}

/// Recursively unpack all mix files in the game directory
/// 
/// # Arguments 
/// 
/// * `folder`: The game directory path
/// * `db`: The global mix database
///
/// # Examples 
/// 
/// ```no_run
/// # use std::path::Path;
/// # use ra2_mix::{decompress, MixDatabase, Ra2Error};
/// fn main() -> Result<(), Ra2Error> {
///     let db = MixDatabase::load(Path::new("C:\\Program Files (x86)\\XCC\\Utilities\\global mix database.dat"))?;
///     decompress(Path::new("C:\\Red Alert 2 - Yuri's Revenge"), &db)?;
///     Ok(())
/// }
/// ```
pub fn decompress(folder: &Path, db: &MixDatabase) -> Result<(), Ra2Error> {
    if !folder.exists() {
        return Err(Ra2Error::FileNotFound("file not found".to_string()));
    }
    if !folder.is_dir() {
        return Err(Ra2Error::FileNotFound("file exists but not folder".to_string()));
    }
    let mut queue = VecDeque::new();
    task_append(&mut queue, folder)?;
    while let Some(path) = queue.pop_front() {
        let mix = MixPackage::load(&path, db)?;
        let sub_folder = path.with_extension("");
        create_dir_all(&sub_folder)?;
        unsafe { mix.dump_unchecked(&sub_folder)? };
        task_append(&mut queue, &sub_folder)?;
    }
    Ok(())
}

fn task_append(paths: &mut VecDeque<PathBuf>, folder: &Path) -> Result<(), Ra2Error> {
    for entry in folder.read_dir()? {
        let entry = entry?;
        let path = entry.path();
        match path.extension() {
            Some(s) if s.eq("mix") => {
                println!("Found mix file: {}", path.display());
                paths.push_back(path);
            }
            _ => {}
        }
    }
    Ok(())
}