1use anyhow::{Context, Result};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::Path;
5use std::process::Command;
6use std::sync::Arc;
7
8use exif::{In, Reader as ExifReader, Tag};
9use hex;
10use image::{self, GenericImageView};
11use img_hash::{HasherConfig, ImageHash};
12use infer;
13use log;
14use mime_guess::MimeGuess;
15
16use crate::audio_fingerprint;
17use crate::file_utils::{DuplicateSet, FileInfo};
18use crate::video_fingerprint;
19
20pub fn is_ffmpeg_available() -> bool {
22 Command::new("ffmpeg").arg("-version").output().is_ok()
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27pub enum MediaKind {
28 Image,
29 Video,
30 Audio,
31 Unknown,
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36pub enum ResolutionPreference {
37 Highest,
38 Lowest,
39 ClosestTo(u32, u32), }
41
42impl std::fmt::Display for ResolutionPreference {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 match self {
45 Self::Highest => write!(f, "highest"),
46 Self::Lowest => write!(f, "lowest"),
47 Self::ClosestTo(w, h) => write!(f, "closest to {}x{}", w, h),
48 }
49 }
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54pub struct FormatPreference {
55 pub formats: Vec<String>, }
57
58impl Default for FormatPreference {
59 fn default() -> Self {
60 Self {
61 formats: vec![
62 "raw".to_string(),
64 "arw".to_string(),
65 "cr2".to_string(),
66 "nef".to_string(),
67 "orf".to_string(),
68 "rw2".to_string(),
69 "png".to_string(),
71 "tiff".to_string(),
72 "bmp".to_string(),
73 "jpg".to_string(),
75 "jpeg".to_string(),
76 "mp4".to_string(),
77 "mov".to_string(),
78 "mp3".to_string(),
79 "flac".to_string(),
80 "wav".to_string(),
81 ],
82 }
83 }
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct MediaDedupOptions {
89 pub enabled: bool,
90 pub resolution_preference: ResolutionPreference,
91 pub format_preference: FormatPreference,
92 pub similarity_threshold: u32, }
94
95impl Default for MediaDedupOptions {
96 fn default() -> Self {
97 Self {
98 enabled: false,
99 resolution_preference: ResolutionPreference::Highest,
100 format_preference: FormatPreference::default(),
101 similarity_threshold: 90, }
103 }
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct MediaMetadata {
109 pub kind: MediaKind,
110 pub width: Option<u32>,
111 pub height: Option<u32>,
112 pub format: String,
113 pub duration: Option<f64>, pub bitrate: Option<u32>,
115 pub perceptual_hash: Option<String>,
116 pub fingerprint: Option<Vec<u8>>,
117}
118
119#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct MediaFileInfo {
122 pub file_info: FileInfo,
123 pub metadata: Option<MediaMetadata>,
124}
125
126impl From<FileInfo> for MediaFileInfo {
127 fn from(file_info: FileInfo) -> Self {
128 Self {
129 file_info,
130 metadata: None,
131 }
132 }
133}
134
135pub fn detect_media_type(path: &Path) -> MediaKind {
137 if let Ok(content) = std::fs::read(path) {
139 if let Some(info) = infer::get(&content) {
140 match info.mime_type() {
141 m if m.starts_with("image/") => return MediaKind::Image,
142 m if m.starts_with("video/") => return MediaKind::Video,
143 m if m.starts_with("audio/") => return MediaKind::Audio,
144 _ => {}
145 }
146 }
147 }
148
149 if let Some(extension) = path.extension() {
151 if let Some(ext_str) = extension.to_str() {
152 let mime = MimeGuess::from_ext(ext_str).first_or_octet_stream();
153 let type_str = mime.type_().as_str();
154
155 if type_str.starts_with("image") {
156 return MediaKind::Image;
157 } else if type_str.starts_with("video") {
158 return MediaKind::Video;
159 } else if type_str.starts_with("audio") {
160 return MediaKind::Audio;
161 }
162 }
163 }
164
165 MediaKind::Unknown
166}
167
168pub fn extract_image_metadata(path: &Path) -> Result<MediaMetadata> {
170 let format = path
171 .extension()
172 .and_then(|e| e.to_str())
173 .unwrap_or("unknown")
174 .to_lowercase();
175
176 let img = image::open(path).with_context(|| format!("Failed to open image: {:?}", path))?;
178
179 let (width, height) = img.dimensions();
180
181 let hasher = HasherConfig::new().to_hasher();
183
184 let img_hash_img = {
187 let rgba8 = img.to_rgba8();
188 let width = rgba8.width();
189 let height = rgba8.height();
190 let raw_pixels = rgba8.into_raw();
191
192 let buffer = img_hash::image::ImageBuffer::from_raw(width, height, raw_pixels)
194 .expect("Failed to convert image buffer");
195
196 img_hash::image::DynamicImage::ImageRgba8(buffer)
198 };
199
200 let hash = hasher.hash_image(&img_hash_img);
202 let hash_str = hex::encode(hash.as_bytes());
203
204 let _bitrate: Option<u32> = None;
206 if let Ok(file) = std::fs::File::open(path) {
207 if let Ok(exif) = ExifReader::new().read_from_container(&mut std::io::BufReader::new(&file))
208 {
209 if let Some(field) = exif.get_field(Tag::XResolution, In::PRIMARY) {
211 if let Some(width) = field.value.get_uint(0) {
213 log::debug!("Image resolution: {}", width);
216 }
217 }
218 }
219 }
220
221 Ok(MediaMetadata {
222 kind: MediaKind::Image,
223 width: Some(width),
224 height: Some(height),
225 format,
226 duration: None, bitrate: None,
228 perceptual_hash: Some(hash_str),
229 fingerprint: None, })
231}
232
233pub fn extract_video_metadata(path: &Path) -> Result<MediaMetadata> {
235 let format = path
236 .extension()
237 .and_then(|e| e.to_str())
238 .unwrap_or("unknown")
239 .to_lowercase();
240
241 if !is_ffmpeg_available() {
243 return Err(anyhow::anyhow!(
244 "ffmpeg is required for video processing but is not available"
245 ));
246 }
247
248 let (width, height, duration, bitrate) = video_fingerprint::extract_video_metadata(path)?;
250
251 let fingerprint = video_fingerprint::fingerprint_video(path)?;
253
254 Ok(MediaMetadata {
255 kind: MediaKind::Video,
256 width,
257 height,
258 format,
259 duration,
260 bitrate,
261 perceptual_hash: None,
262 fingerprint: Some(fingerprint),
263 })
264}
265
266pub fn extract_audio_metadata(path: &Path) -> Result<MediaMetadata> {
268 let format = path
269 .extension()
270 .and_then(|e| e.to_str())
271 .unwrap_or("unknown")
272 .to_lowercase();
273
274 if !is_ffmpeg_available() {
276 let chromaprint_available = Command::new("fpcalc").arg("-version").output().is_ok();
278
279 if !chromaprint_available {
280 return Err(anyhow::anyhow!(
281 "Neither ffmpeg nor chromaprint is available for audio processing"
282 ));
283 }
284 }
285
286 let fingerprint = audio_fingerprint::fingerprint_file(path)?;
288
289 let mut duration = None;
291 let mut bitrate = None;
292
293 if is_ffmpeg_available() {
294 let output = Command::new("ffprobe")
296 .args([
297 "-v",
298 "error",
299 "-select_streams",
300 "a:0", "-show_entries",
302 "stream=duration,bit_rate",
303 "-of",
304 "json",
305 path.to_str().unwrap(),
306 ])
307 .output();
308
309 if let Ok(output) = output {
310 if output.status.success() {
311 let output_str = String::from_utf8_lossy(&output.stdout);
312 if let Ok(json) = serde_json::from_str::<serde_json::Value>(&output_str) {
313 duration = json["streams"][0]["duration"]
315 .as_str()
316 .and_then(|d| d.parse::<f64>().ok())
317 .or_else(|| json["streams"][0]["duration"].as_f64());
318
319 bitrate = json["streams"][0]["bit_rate"]
321 .as_str()
322 .and_then(|b| b.parse::<u32>().ok())
323 .or_else(|| json["streams"][0]["bit_rate"].as_u64().map(|b| b as u32));
324 }
325 }
326 }
327 }
328
329 Ok(MediaMetadata {
330 kind: MediaKind::Audio,
331 width: None, height: None, format,
334 duration,
335 bitrate,
336 perceptual_hash: None,
337 fingerprint: Some(fingerprint),
338 })
339}
340
341pub fn extract_media_metadata(path: &Path) -> Result<MediaMetadata> {
343 let media_kind = detect_media_type(path);
344
345 match media_kind {
346 MediaKind::Image => extract_image_metadata(path),
347 MediaKind::Video => extract_video_metadata(path),
348 MediaKind::Audio => extract_audio_metadata(path),
349 MediaKind::Unknown => Err(anyhow::anyhow!("Unknown media type for path: {:?}", path)),
350 }
351}
352
353pub fn calculate_image_similarity(hash1: &str, hash2: &str) -> u32 {
355 let parse_hash = |hash_str: &str| -> Option<ImageHash> {
357 let bytes = (0..hash_str.len())
358 .step_by(2)
359 .map(|i| u8::from_str_radix(&hash_str[i..i + 2], 16).ok())
360 .collect::<Option<Vec<u8>>>()?;
361
362 ImageHash::from_bytes(&bytes).ok()
363 };
364
365 if let (Some(img_hash1), Some(img_hash2)) = (parse_hash(hash1), parse_hash(hash2)) {
366 let distance = img_hash1.dist(&img_hash2);
368
369 let max_distance = 64; let similarity = ((max_distance - distance) as f64 / max_distance as f64) * 100.0;
372
373 return similarity as u32;
374 }
375
376 0 }
378
379pub fn calculate_video_similarity(fp1: &[u8], fp2: &[u8]) -> u32 {
381 (video_fingerprint::compare_fingerprints(fp1, fp2) * 100.0) as u32
383}
384
385pub fn calculate_audio_similarity(fp1: &[u8], fp2: &[u8]) -> u32 {
387 (audio_fingerprint::compare_fingerprints(fp1, fp2) * 100.0) as u32
389}
390
391pub fn compare_media_files(a: &MediaFileInfo, b: &MediaFileInfo) -> u32 {
393 match (&a.metadata, &b.metadata) {
394 (Some(meta_a), Some(meta_b)) => {
395 if meta_a.kind != meta_b.kind {
397 return 0;
398 }
399
400 match meta_a.kind {
401 MediaKind::Image => match (&meta_a.perceptual_hash, &meta_b.perceptual_hash) {
402 (Some(hash_a), Some(hash_b)) => calculate_image_similarity(hash_a, hash_b),
403 _ => 0,
404 },
405 MediaKind::Video => match (&meta_a.fingerprint, &meta_b.fingerprint) {
406 (Some(fp_a), Some(fp_b)) => calculate_video_similarity(fp_a, fp_b),
407 _ => 0,
408 },
409 MediaKind::Audio => match (&meta_a.fingerprint, &meta_b.fingerprint) {
410 (Some(fp_a), Some(fp_b)) => calculate_audio_similarity(fp_a, fp_b),
411 _ => 0,
412 },
413 MediaKind::Unknown => 0,
414 }
415 }
416 _ => 0, }
418}
419
420pub fn determine_preferred_media_file<'a>(
422 files: &'a [MediaFileInfo],
423 options: &'a MediaDedupOptions,
424) -> Option<&'a MediaFileInfo> {
425 if files.is_empty() {
426 return None;
427 }
428
429 let files_with_metadata: Vec<_> = files.iter().filter(|f| f.metadata.is_some()).collect();
431
432 if files_with_metadata.is_empty() {
433 return files.first(); }
435
436 let format_ranks: HashMap<String, usize> = options
438 .format_preference
439 .formats
440 .iter()
441 .enumerate()
442 .map(|(i, fmt)| (fmt.clone(), i))
443 .collect();
444
445 let get_format_rank = |file: &&MediaFileInfo| -> usize {
447 let format = match &file.metadata {
448 Some(meta) => &meta.format,
449 None => return usize::MAX, };
451 *format_ranks.get(format).unwrap_or(&usize::MAX)
452 };
453
454 let mut preferred_format_files = files_with_metadata.clone();
456 preferred_format_files.sort_by_key(get_format_rank);
457
458 let best_format_rank = get_format_rank(&preferred_format_files[0]);
460 let best_format_files: Vec<_> = preferred_format_files
461 .into_iter()
462 .filter(|f| get_format_rank(f) == best_format_rank)
463 .collect();
464
465 if best_format_files.len() > 1 {
467 match options.resolution_preference {
468 ResolutionPreference::Highest => {
469 best_format_files
471 .into_iter()
472 .max_by_key(|file| match &file.metadata {
473 Some(meta) => meta.width.unwrap_or(0) * meta.height.unwrap_or(0),
474 None => 0,
475 })
476 }
477 ResolutionPreference::Lowest => {
478 best_format_files
480 .into_iter()
481 .min_by_key(|file| match &file.metadata {
482 Some(meta) => {
483 meta.width.unwrap_or(u32::MAX) * meta.height.unwrap_or(u32::MAX)
484 }
485 None => u32::MAX,
486 })
487 }
488 ResolutionPreference::ClosestTo(target_width, target_height) => {
489 best_format_files.into_iter().min_by_key(|file| {
491 match &file.metadata {
492 Some(meta) => {
493 let w = meta.width.unwrap_or(0);
494 let h = meta.height.unwrap_or(0);
495 let dw = if w > target_width {
496 w - target_width
497 } else {
498 target_width - w
499 };
500 let dh = if h > target_height {
501 h - target_height
502 } else {
503 target_height - h
504 };
505 dw * dw + dh * dh }
507 None => u32::MAX,
508 }
509 })
510 }
511 }
512 } else {
513 best_format_files.into_iter().next()
515 }
516}
517
518#[allow(clippy::arc_with_non_send_sync)]
519pub fn find_similar_media_files(
520 file_infos: &[FileInfo],
521 options: &MediaDedupOptions,
522 progress_callback: Option<Box<dyn Fn(usize, usize) + Send>>,
523) -> Result<Vec<Vec<MediaFileInfo>>> {
524 if !options.enabled {
525 return Ok(Vec::new());
526 }
527
528 log::info!(
529 "Starting media deduplication with threshold: {}%",
530 options.similarity_threshold
531 );
532
533 let has_video_files = file_infos.iter().any(|f| {
535 let kind = detect_media_type(&f.path);
536 kind == MediaKind::Video
537 });
538
539 if has_video_files && !is_ffmpeg_available() {
540 log::warn!("FFmpeg is not installed. Video deduplication will be limited.");
541 }
542
543 let total_files = file_infos.len();
545 let mut processed = 0;
546
547 let progress_callback = progress_callback.map(|cb| {
549 let cb = Arc::new(cb);
550 move |count, total| {
551 let cb = cb.clone();
552 cb(count, total);
553 }
554 });
555
556 let media_files: Vec<MediaFileInfo> = file_infos
558 .iter()
559 .map(|file_info| {
560 let mut media_file = MediaFileInfo::from(file_info.clone());
561
562 let media_kind = detect_media_type(&file_info.path);
564 if media_kind != MediaKind::Unknown {
565 media_file.metadata = match extract_media_metadata(&file_info.path) {
566 Ok(metadata) => Some(metadata),
567 Err(e) => {
568 log::warn!(
569 "Failed to extract media metadata for {:?}: {}",
570 file_info.path,
571 e
572 );
573 None
574 }
575 };
576 }
577
578 processed += 1;
580 if let Some(cb) = &progress_callback {
581 cb(processed, total_files);
582 }
583
584 media_file
585 })
586 .filter(|media_file| media_file.metadata.is_some())
587 .collect();
588
589 log::info!("Extracted metadata for {} media files", media_files.len());
590
591 let mut image_files: Vec<_> = Vec::new();
593 let mut video_files: Vec<_> = Vec::new();
594 let mut audio_files: Vec<_> = Vec::new();
595
596 for file in &media_files {
597 if let Some(metadata) = &file.metadata {
598 match metadata.kind {
599 MediaKind::Image => image_files.push(file),
600 MediaKind::Video => video_files.push(file),
601 MediaKind::Audio => audio_files.push(file),
602 _ => {}
603 }
604 }
605 }
606
607 log::info!(
608 "Media file count: {} images, {} videos, {} audio files",
609 image_files.len(),
610 video_files.len(),
611 audio_files.len()
612 );
613
614 let mut similar_groups: Vec<Vec<MediaFileInfo>> = Vec::new();
616
617 process_media_type_similarity(&image_files, options, &mut similar_groups)?;
619 process_media_type_similarity(&video_files, options, &mut similar_groups)?;
620 process_media_type_similarity(&audio_files, options, &mut similar_groups)?;
621
622 log::info!(
623 "Found {} groups of similar media files.",
624 similar_groups.len()
625 );
626
627 Ok(similar_groups)
628}
629
630pub fn process_media_type_similarity(
632 files: &[&MediaFileInfo],
633 options: &MediaDedupOptions,
634 similar_groups: &mut Vec<Vec<MediaFileInfo>>,
635) -> Result<()> {
636 if files.len() < 2 {
637 return Ok(());
638 }
639
640 let mut processed = vec![false; files.len()];
642
643 for i in 0..files.len() {
645 if processed[i] {
646 continue;
647 }
648
649 let mut current_group = Vec::new();
650 current_group.push(files[i].clone());
651 processed[i] = true;
652
653 for j in i + 1..files.len() {
654 if processed[j] {
655 continue;
656 }
657
658 let similarity = compare_media_files(files[i], files[j]);
659 if similarity >= options.similarity_threshold {
660 current_group.push(files[j].clone());
661 processed[j] = true;
662 }
663 }
664
665 if current_group.len() > 1 {
666 similar_groups.push(current_group);
667 }
668 }
669
670 Ok(())
671}
672
673pub fn convert_to_duplicate_sets(
675 similar_groups: &[Vec<MediaFileInfo>],
676 options: &MediaDedupOptions,
677) -> Vec<DuplicateSet> {
678 let mut duplicate_sets = Vec::new();
679
680 for group in similar_groups {
681 if group.len() < 2 {
682 continue;
683 }
684
685 let kept_file = determine_preferred_media_file(group, options);
687
688 if let Some(kept) = kept_file {
689 let mut file_infos = group
691 .iter()
692 .map(|mf| mf.file_info.clone())
693 .collect::<Vec<_>>();
694
695 if let Some(kept_idx) = file_infos
697 .iter()
698 .position(|f| f.path == kept.file_info.path)
699 {
700 let kept_file_info = file_infos.remove(kept_idx);
701 file_infos.insert(0, kept_file_info);
702 }
703
704 let hash = format!("media_{}", group[0].file_info.path.to_string_lossy());
706 let size = group[0].file_info.size;
707
708 duplicate_sets.push(DuplicateSet {
709 files: file_infos,
710 size,
711 hash,
712 });
713 }
714 }
715
716 duplicate_sets
717}
718
719pub fn add_media_options_to_cli(
721 options: &mut MediaDedupOptions,
722 enable: bool,
723 resolution: &str,
724 formats: &[String],
725 threshold: u32,
726) {
727 options.enabled = enable;
728
729 match resolution {
731 "highest" => options.resolution_preference = ResolutionPreference::Highest,
732 "lowest" => options.resolution_preference = ResolutionPreference::Lowest,
733 custom => {
734 if let Some((width, height)) = custom.split_once('x') {
736 if let (Ok(w), Ok(h)) = (width.parse::<u32>(), height.parse::<u32>()) {
737 options.resolution_preference = ResolutionPreference::ClosestTo(w, h);
738 }
739 }
740 }
741 }
742
743 if !formats.is_empty() {
745 options.format_preference.formats = formats.to_vec();
746 }
747
748 if threshold > 0 && threshold <= 100 {
750 options.similarity_threshold = threshold;
751 }
752}
753
754#[cfg(test)]
755mod tests {
756 use super::*;
757 use std::path::PathBuf;
758 use std::time::SystemTime;
759
760 fn create_test_file_info(path: &str, size: u64) -> FileInfo {
762 FileInfo {
763 path: PathBuf::from(path),
764 size,
765 hash: Some("test_hash".to_string()),
766 modified_at: Some(SystemTime::now()),
767 created_at: Some(SystemTime::now()),
768 }
769 }
770
771 #[test]
772 fn test_ffmpeg_availability() {
773 let available = is_ffmpeg_available();
775 println!("FFmpeg available: {}", available);
777 }
778
779 #[test]
780 fn test_media_kind_from_extension() {
781 assert_eq!(detect_media_type(Path::new("test.jpg")), MediaKind::Image);
783 assert_eq!(detect_media_type(Path::new("test.png")), MediaKind::Image);
784
785 assert_eq!(detect_media_type(Path::new("test.mp4")), MediaKind::Video);
787 assert_eq!(detect_media_type(Path::new("test.mov")), MediaKind::Video);
788
789 assert_eq!(detect_media_type(Path::new("test.mp3")), MediaKind::Audio);
791 assert_eq!(detect_media_type(Path::new("test.wav")), MediaKind::Audio);
792
793 assert_eq!(detect_media_type(Path::new("test.txt")), MediaKind::Unknown);
795 }
796
797 #[test]
798 fn test_format_preference() {
799 let format_pref = FormatPreference::default();
801
802 assert!(
804 format_pref.formats.iter().position(|f| f == "raw").unwrap()
805 < format_pref.formats.iter().position(|f| f == "jpg").unwrap()
806 );
807
808 assert!(
810 format_pref.formats.iter().position(|f| f == "png").unwrap()
811 < format_pref.formats.iter().position(|f| f == "jpg").unwrap()
812 );
813 }
814
815 #[test]
816 fn test_resolution_preference_display() {
817 assert_eq!(ResolutionPreference::Highest.to_string(), "highest");
818 assert_eq!(ResolutionPreference::Lowest.to_string(), "lowest");
819 assert_eq!(
820 ResolutionPreference::ClosestTo(1280, 720).to_string(),
821 "closest to 1280x720"
822 );
823 }
824
825 #[test]
826 fn test_media_dedup_options_default() {
827 let options = MediaDedupOptions::default();
828 assert!(!options.enabled);
829 assert_eq!(options.similarity_threshold, 90);
830
831 match options.resolution_preference {
833 ResolutionPreference::Highest => (),
834 _ => panic!("Default resolution preference should be Highest"),
835 }
836 }
837}