virittaa

file_io.rs (4979B)

---

 use anyhow::Result;
 use lofty::config::WriteOptions;
 use lofty::prelude::*;
 use lofty::probe::Probe;
 use lofty::tag::{ItemKey, Tag};
 use std::path::Path;
 
 use crate::audio_processing::{calculate_bpm, calculate_key};
 use crate::types::TrackInfo;
 use bliss_audio::decoder::Decoder as DecoderTrait;
 use bliss_audio::decoder::ffmpeg::FFmpegDecoder as Decoder;
 
 #[derive(Debug, Default)]
 pub struct FileMetadata {
     pub artist: Option<String>,
     pub track: Option<String>,
 }
 
 fn build_track_info(path: &Path, metadata: FileMetadata, bpm: f32, key: String) -> TrackInfo {
     TrackInfo {
         path: path.to_path_buf(),
         artist: metadata.artist,
         track: metadata.track,
         bpm,
         key,
     }
 }
 
 /// Process a single file, detecting BPM and key.
 ///
 /// # Errors
 ///
 /// Returns an error string if the file cannot be decoded or key detection fails.
 pub fn process_file(path: &Path, write_tags: bool, force: bool) -> Result<TrackInfo, String> {
     let (metadata, existing_tags) = read_tags(path);
 
     if !force && let Some((bpm, key)) = existing_tags {
         println!(
             "File: {}, BPM: {:.1}, Key: {} (existing tags)",
             path.display(),
             bpm,
             key
         );
         return Ok(build_track_info(path, metadata, bpm, key));
     }
 
     let song = Decoder::decode(path).map_err(|e| format!("Error decoding: {e}"))?;
     let samples = &song.sample_array;
 
     let bpm = calculate_bpm(samples);
     let key = calculate_key(samples).map_err(|e| format!("Error calculating key: {e}"))?;
 
     println!("File: {}, BPM: {bpm:.1}, Key: {key}", path.display());
 
     if write_tags && let Err(e) = write_metadata(path, bpm, &key) {
         eprintln!("Error writing metadata for {}: {e}", path.display());
     }
 
     Ok(build_track_info(path, metadata, bpm, key))
 }
 
 fn read_tags(path: &Path) -> (FileMetadata, Option<(f32, String)>) {
     match Probe::open(path).and_then(Probe::read) {
         Ok(tagged_file) => {
             let tag = tagged_file
                 .primary_tag()
                 .or_else(|| tagged_file.first_tag());
             let Some(t) = tag else {
                 return (FileMetadata::default(), None);
             };
 
             let metadata = FileMetadata {
                 artist: t
                     .get_string(ItemKey::TrackArtist)
                     .or_else(|| t.get_string(ItemKey::AlbumArtist))
                     .map(String::from),
                 track: t.get_string(ItemKey::TrackTitle).map(String::from),
             };
 
             let existing_tags = t
                 .get_string(ItemKey::Bpm)
                 .zip(t.get_string(ItemKey::InitialKey))
                 .and_then(|(bpm_str, key)| {
                     bpm_str
                         .parse::<f32>()
                         .ok()
                         .map(|bpm| (bpm, key.to_string()))
                 });
 
             (metadata, existing_tags)
         }
         Err(_) => (FileMetadata::default(), None),
     }
 }
 
 /// Write BPM and key metadata to a file.
 ///
 /// # Errors
 ///
 /// Returns an error if the file cannot be opened, read, or written.
 ///
 /// # Panics
 ///
 /// Panics if there are no tags and a new tag cannot be created.
 pub fn write_metadata(path: &Path, bpm: f32, key: &str) -> Result<()> {
     let mut tagged_file = Probe::open(path)?.read()?;
     let tag = match tagged_file.primary_tag_mut() {
         Some(primary_tag) => primary_tag,
         None => {
             if let Some(first_tag) = tagged_file.first_tag_mut() {
                 first_tag
             } else {
                 let tag_type = tagged_file.primary_tag_type();
                 tagged_file.insert_tag(Tag::new(tag_type));
                 tagged_file
                     .primary_tag_mut()
                     .expect("Failed to create new tag")
             }
         }
     };
 
     tag.insert_text(ItemKey::Bpm, bpm.round().to_string());
     tag.insert_text(ItemKey::InitialKey, key.to_string());
     if let Some(tkey) = ItemKey::from_key(tag.tag_type(), "TKEY") {
         tag.insert_text(tkey, key.to_string());
     }
 
     tag.save_to_path(path, WriteOptions::default())?;
     Ok(())
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn build_track_info_carries_metadata_and_analysis() {
         let metadata = FileMetadata {
             artist: Some("Aphex Twin".to_string()),
             track: Some("Xtal".to_string()),
         };
         let info = build_track_info(Path::new("/music/xtal.flac"), metadata, 120.5, "A Minor".to_string());
 
         assert_eq!(info.path, Path::new("/music/xtal.flac"));
         assert_eq!(info.artist.as_deref(), Some("Aphex Twin"));
         assert_eq!(info.track.as_deref(), Some("Xtal"));
         assert!((info.bpm - 120.5).abs() < f32::EPSILON);
         assert_eq!(info.key, "A Minor");
     }
 
     #[test]
     fn file_metadata_defaults_to_none() {
         let md = FileMetadata::default();
         assert!(md.artist.is_none());
         assert!(md.track.is_none());
     }
 }