use anyhow::Result;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use zip::ZipArchive;
type ImageList<'a> = Vec<(&'a String, &'a PathBuf)>;
type ExtractedImages = Vec<(String, PathBuf)>;
#[derive(Debug)]
pub struct ImageExtractor {
temp_dir: PathBuf,
extracted_images: HashMap<String, PathBuf>, }
#[derive(Debug, Clone)]
pub enum ImageFormat {
Png,
Jpeg,
Gif,
Bmp,
Tiff,
}
impl ImageFormat {
pub fn from_filename(filename: &str) -> Option<Self> {
let extension = Path::new(filename).extension()?.to_str()?.to_lowercase();
match extension.as_str() {
"png" => Some(Self::Png),
"jpg" | "jpeg" => Some(Self::Jpeg),
"gif" => Some(Self::Gif),
"bmp" => Some(Self::Bmp),
"tiff" | "tif" => Some(Self::Tiff),
_ => None,
}
}
pub fn to_extension(&self) -> &'static str {
match self {
Self::Png => "png",
Self::Jpeg => "jpg",
Self::Gif => "gif",
Self::Bmp => "bmp",
Self::Tiff => "tiff",
}
}
}
impl ImageExtractor {
pub fn new() -> Result<Self> {
let temp_dir = std::env::temp_dir().join("doxx_images");
fs::create_dir_all(&temp_dir)?;
Ok(Self {
temp_dir,
extracted_images: HashMap::new(),
})
}
pub fn extract_images_from_docx(&mut self, docx_path: &Path) -> Result<()> {
let file = File::open(docx_path)?;
let mut archive = ZipArchive::new(file)?;
for i in 0..archive.len() {
let mut file = archive.by_index(i)?;
let outpath = file.name().to_string();
if outpath.starts_with("word/media/") && self.is_image_file(&outpath) {
let filename = Path::new(&outpath)
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown");
let temp_file_path = self.temp_dir.join(filename);
let mut buffer = Vec::new();
file.read_to_end(&mut buffer)?;
let mut temp_file = File::create(&temp_file_path)?;
temp_file.write_all(&buffer)?;
let rel_id = filename.to_string(); self.extracted_images.insert(rel_id, temp_file_path);
}
}
println!(
"Extracted {} images to {}",
self.extracted_images.len(),
self.temp_dir.display()
);
Ok(())
}
pub fn get_image_data(&self, rel_id: &str) -> Result<Vec<u8>> {
if let Some(path) = self.extracted_images.get(rel_id) {
Ok(fs::read(path)?)
} else {
anyhow::bail!("Image not found: {}", rel_id)
}
}
pub fn get_image_path(&self, rel_id: &str) -> Option<&PathBuf> {
self.extracted_images.get(rel_id)
}
pub fn list_images(&self) -> ImageList<'_> {
self.extracted_images.iter().collect()
}
pub fn get_extracted_images(&self) -> ExtractedImages {
self.extracted_images
.iter()
.map(|(rel_id, path)| (rel_id.clone(), path.clone()))
.collect()
}
pub fn get_extracted_images_sorted(&self) -> ExtractedImages {
let mut images: ExtractedImages = self
.extracted_images
.iter()
.map(|(rel_id, path)| (rel_id.clone(), path.clone()))
.collect();
images.sort_by(|a, b| a.0.cmp(&b.0));
images
}
pub fn cleanup(&self) -> Result<()> {
if self.temp_dir.exists() {
fs::remove_dir_all(&self.temp_dir)?;
}
Ok(())
}
fn is_image_file(&self, filename: &str) -> bool {
ImageFormat::from_filename(filename).is_some()
}
}
impl Drop for ImageExtractor {
fn drop(&mut self) {
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_image_format_detection() {
assert!(matches!(
ImageFormat::from_filename("image.png"),
Some(ImageFormat::Png)
));
assert!(matches!(
ImageFormat::from_filename("photo.jpg"),
Some(ImageFormat::Jpeg)
));
assert!(matches!(
ImageFormat::from_filename("photo.jpeg"),
Some(ImageFormat::Jpeg)
));
assert!(matches!(
ImageFormat::from_filename("animation.gif"),
Some(ImageFormat::Gif)
));
assert!(matches!(
ImageFormat::from_filename("bitmap.bmp"),
Some(ImageFormat::Bmp)
));
assert!(matches!(
ImageFormat::from_filename("scan.tiff"),
Some(ImageFormat::Tiff)
));
assert!(ImageFormat::from_filename("document.txt").is_none());
}
#[test]
fn test_image_extractor_creation() {
let extractor = ImageExtractor::new().unwrap();
assert!(extractor.temp_dir.exists());
assert!(extractor.extracted_images.is_empty());
}
}