cbz2pdf/src/formats/cbz.rs

use std::ffi::OsStr;
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;

use anyhow::Result;
use zip::ZipArchive;

use crate::model::Document;

use super::{CbxReader, FormatReader, FormatWriter};

pub struct CbzReader;

impl CbxReader for CbzReader {
    fn extract_images(&self, input: &Path) -> Result<Vec<(String, Vec<u8>)>> {
        let mut zip = ZipArchive::new(File::open(input)?)?;
        let mut files: Vec<(String, Vec<u8>)> = Vec::new();
        for i in 0..zip.len() {
            let mut file = zip.by_index(i)?;
            let mut image_data = Vec::new();
            let name = file
                .enclosed_name()
                .expect("Failed to read file name")
                .to_owned();
            if name.extension() == Some(OsStr::new("jpg")) {
                file.read_to_end(&mut image_data)?;
                files.push((
                    name.file_name()
                        .expect("Failed to read file name")
                        .to_string_lossy()
                        .to_string(),
                    image_data,
                ));
            }
        }
        Ok(files)
    }
}

impl FormatReader for CbzReader {
    fn read(&self, input: &Path) -> Result<Document> {
        self.read_cbx(input)
    }
}

pub struct CbzWriter;

impl FormatWriter for CbzWriter {
    fn write(&self, doc: &Document, output: &Path) -> Result<()> {
        use zip::write::SimpleFileOptions;
        let file = File::create(output)?;
        let mut zip = zip::ZipWriter::new(file);
        let options = SimpleFileOptions::default();

        for (idx, page) in doc.pages.iter().enumerate() {
            let mut name = page.name.clone();
            if Path::new(&name).extension().and_then(OsStr::to_str) != Some("jpg") {
                name = format!("{:03}.jpg", idx + 1);
            }
            zip.start_file(&name, options)?;
            if let Some(dct) = &page.jpeg_dct {
                zip.write_all(dct)?;
            } else {
                // Encode to JPEG
                let rgb = page.image.to_rgb8();
                let (w, h) = (rgb.width(), rgb.height());
                let mut cursor = std::io::Cursor::new(Vec::new());
                {
                    let mut enc =
                        image::codecs::jpeg::JpegEncoder::new_with_quality(&mut cursor, 85);
                    enc.encode(&rgb.into_raw(), w, h, image::ColorType::Rgb8.into())?;
                }
                let data = cursor.into_inner();
                zip.write_all(&data)?;
            }
        }

        zip.finish()?;
        Ok(())
    }
}