msgmerge/src/types/project.rs

use std::{collections::HashSet, io::{self, Cursor, Read, Write}, path::PathBuf, process::exit};

use serde::{Deserialize, Serialize};
use tokio::{fs::File, io::{AsyncReadExt, AsyncWriteExt, BufWriter}};
use zip::{write::FileOptions, ZipArchive, ZipWriter};

use super::{message::Message, user::User};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Project {
    pub users: Vec<User>,
    pub messages: Vec<Message>,
    pub timestamps: HashSet<usize>,
    pub media_index: usize,
}

pub struct ProjectRaw {
    pub project: Project,
    pub zip: ZipArchive<Cursor<Vec<u8>>>
}

impl Project {
    pub async fn load(pathstr: String) -> io::Result<ProjectRaw> {
        let path = PathBuf::from(pathstr);

        if !path.exists() {
            return Err(io::Error::new(io::ErrorKind::NotFound, "File does not exist!"));
        }

        let mut file = File::open(path).await?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer).await?;

        let cursor = Cursor::new(buffer);
        let archive = ZipArchive::new(cursor)?;

        let mut copy_archive = archive.clone();
        let mut project_file = copy_archive.by_name("project.json").unwrap_or_else(|_| {
            eprintln!("Invalid project format!");
            exit(1);
        });

        let mut json_content = String::new();
        project_file.read_to_string(&mut json_content)?;

        let data: Project = serde_json::from_str(&json_content)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        Ok(ProjectRaw {
            project: data,
            zip: archive
        })
    }

    pub async fn save_new(&self, file_path: String) -> io::Result<()> {
        let contents = serde_json::to_string(&self)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        let file = File::create(file_path).await?;
        let mut writer = BufWriter::new(file);

        let mut buffer = Cursor::new(Vec::new());

        {
            let mut zip = ZipWriter::new(&mut buffer);

            let options = FileOptions::default().unix_permissions(0o755);
            zip.add_directory("media/", options)?;

            zip.start_file("project.json", options)?;
            zip.write_all(contents.as_bytes())?;

            zip.finish()?;
        }

        writer.write_all(&buffer.into_inner()).await?;
        writer.flush().await?;

        Ok(())
    }

    pub fn push_msg(&mut self, message: &Message) -> bool {
        if self.timestamps.contains(&message.timestamp) {
            return false;
        }

        self.timestamps.insert(message.timestamp);
        self.messages.push(message.clone());

        return true;
    }

    pub fn sort_messages(&mut self) {
        self.messages.sort_by_key(|msg| msg.timestamp);
    }
}


impl ProjectRaw {
    /// Save multiple files to the zip archive.
    ///
    /// # Arguments
    /// * `files` - A list of tuples where each tuple contains:
    ///     - `target_path`: The path inside the zip archive (e.g., `media/1`).
    ///     - `source_path`: The path to the file on the filesystem.
    pub async fn save_media_files(
        &mut self,
        files: Vec<(String, String)>,
    ) -> io::Result<()> {
        let mut new_buffer = Cursor::new(Vec::new());
        {
            let mut zip_writer = ZipWriter::new(&mut new_buffer);
            let options = FileOptions::default().unix_permissions(0o644);

            let replace_files: HashSet<_> = files.iter().map(|(target, _)| target.clone()).collect();

            for i in 0..self.zip.len() {
                let mut file = self.zip.by_index(i)?;
                let name = file.name().to_string();

                if replace_files.contains(&name) {
                    continue;
                }

                zip_writer.start_file(name, options)?;
                let mut file_content = Vec::new();
                file.read_to_end(&mut file_content)?;
                zip_writer.write_all(&file_content)?;
            }

            for (target_path, source_path) in files {
                let mut file = File::open(source_path).await?;
                let mut file_content = Vec::new();
                file.read_to_end(&mut file_content).await?;

                zip_writer.start_file(target_path, options)?;
                zip_writer.write_all(&file_content)?;
            }

            zip_writer.finish()?;
        }

        self.zip = ZipArchive::new(new_buffer)?;

        Ok(())
    }

    /// Save the `Project` to `project.json` and flush the zip archive to the specified file.
    ///
    /// # Arguments
    /// * `output_path` - The path to save the zip archive to.
    pub async fn save(&mut self, output_path: &str) -> io::Result<()> {
        let project_json = serde_json::to_string(&self.project)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        let mut new_buffer = Cursor::new(Vec::new());
        {
            let mut zip_writer = ZipWriter::new(&mut new_buffer);
            let options = FileOptions::default().unix_permissions(0o644);

            let mut project_json_written = false;

            for i in 0..self.zip.len() {
                let mut file = self.zip.by_index(i)?;
                let name = file.name();

                if name == "project.json" {
                    // Replace the existing `project.json`.
                    zip_writer.start_file("project.json", options)?;
                    zip_writer.write_all(project_json.as_bytes())?;
                    project_json_written = true;
                } else {
                    // Copy other files as-is.
                    zip_writer.start_file(name, options)?;
                    let mut file_content = Vec::new();
                    file.read_to_end(&mut file_content)?;
                    zip_writer.write_all(&file_content)?;
                }
            }

            if !project_json_written {
                zip_writer.start_file("project.json", options)?;
                zip_writer.write_all(project_json.as_bytes())?;
            }

            zip_writer.finish()?;
        }

        let mut output_file = File::create(output_path).await?;
        output_file.write_all(&new_buffer.into_inner()).await?;
        output_file.flush().await?;

        Ok(())
    }
}