(imperfect) Album art fetching

Signed-off-by: Slendi <slendi@socopon.com>

src/main.rs

@@ -1,21 +1,38 @@
 use std::{
+    collections::{HashMap, HashSet},
     path::PathBuf,
-    thread::sleep,
-    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicBool, Ordering},
+        mpsc,
+    },
+    time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use clap::Parser;
 use discord_presence::{
     Client, Event,
-    models::{ActivityTimestamps, ActivityType},
+    models::{ActivityAssets, ActivityTimestamps, ActivityType},
 };
-use mpd::{Idle, Song, State, Subsystem};
+use mpd::Idle as _;
+use mpd::{Song, State, Subsystem, song::Id};
+use musicbrainz_rs::{
+    FetchCoverart,
+    api::search_query::Search,
+    entity::{
+        recording::{Recording, RecordingSearchQuery},
+        release::ReleaseSearchQuery,
+    },
+};
+use reqwest::blocking::Client as HttpClient;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Serialize, Deserialize)]
 struct Config {
     address: String,
     client_id: u64,
+    small_image: Option<String>,
+    small_text: Option<String>,
 }
 
 impl Default for Config {
@@ -23,6 +40,8 @@ impl Default for Config {
         Self {
             address: "localhost:6600".into(),
             client_id: 1464985070992363645,
+            small_image: None,
+            small_text: None,
         }
     }
 }
@@ -34,16 +53,41 @@ struct Args {
 }
 
 struct Application {
-    mpdc: mpd::Client,
+    mpdc: Arc<Mutex<mpd::Client>>,
     client: Client,
     config: Config,
+    album_art_cache: HashMap<String, Option<String>>,
+    album_art_pending: HashSet<String>,
+    update_tx: mpsc::Sender<UpdateMessage>,
+    update_rx: mpsc::Receiver<UpdateMessage>,
+    running: Arc<AtomicBool>,
+}
+
+const MPD_LOGO: &str = "https://www.musicpd.org/logo.png";
+
+#[derive(Debug)]
+enum UpdateMessage {
+    Player,
+    AlbumArt {
+        cache_key: String,
+        cover_url: Option<String>,
+    },
+    Stop,
 }
 
 impl Application {
     pub fn new(config: Config) -> Self {
+        let (update_tx, update_rx) = mpsc::channel();
         let mut ret = Self {
-            mpdc: mpd::Client::connect(config.address.clone()).unwrap(),
+            mpdc: Arc::new(Mutex::new(
+                mpd::Client::connect(config.address.clone()).unwrap(),
+            )),
             client: Client::new(config.client_id),
+            album_art_cache: HashMap::new(),
+            album_art_pending: HashSet::new(),
+            update_tx,
+            update_rx,
+            running: Arc::new(AtomicBool::new(true)),
             config,
         };
 
@@ -53,6 +97,7 @@ impl Application {
 
         ret.client.start();
         ret.client.block_until_event(Event::Ready).unwrap();
+        let _ = ret.mpdc.lock().unwrap().idle(&[Subsystem::Player]);
         assert!(Client::is_ready());
 
         ret
@@ -65,9 +110,224 @@ impl Application {
             .map(|(_, v)| v.clone())
     }
 
-    pub fn run_update(&mut self) {
-        let _ = self.mpdc.idle(&[Subsystem::Player]);
-        let status = self.mpdc.status().unwrap();
+    fn fetch_album_art(&mut self, song: &mpd::Song) -> Option<String> {
+        let album = Self::tag(song, "Album");
+        let title = song.title.clone();
+        let artist = match song.artist.clone() {
+            Some(artist) => artist,
+            None => return self.cache_missing_album_art(song),
+        };
+
+        if album.is_none() && title.is_none() {
+            return self.cache_missing_album_art(song);
+        }
+
+        let tagged_mbid = Self::tag(song, "MUSICBRAINZ_RELEASEID")
+            .or_else(|| Self::tag(song, "MusicBrainz Release Id"))
+            .or_else(|| Self::tag(song, "MusicBrainz Release ID"))
+            .or_else(|| Self::tag(song, "MBID"));
+
+        let cache_key = tagged_mbid
+            .as_ref()
+            .map(|mbid| format!("mbid:{mbid}"))
+            .or_else(|| {
+                album
+                    .as_ref()
+                    .map(|album| format!("artist:{artist}|album:{album}"))
+            })
+            .or_else(|| {
+                title
+                    .as_ref()
+                    .map(|title| format!("artist:{artist}|title:{title}"))
+            })
+            .unwrap_or_else(|| format!("file:{}", song.file));
+
+        if let Some(cached) = self.album_art_cache.get(&cache_key) {
+            return cached.clone();
+        }
+
+        if self.album_art_pending.contains(&cache_key) {
+            return None;
+        }
+
+        self.album_art_pending.insert(cache_key.clone());
+        let album = album.clone();
+        let title = title.clone();
+        let artist = artist.clone();
+        let tx = self.update_tx.clone();
+        std::thread::spawn(move || {
+            let cover_url = Self::fetch_album_art_blocking(
+                album.as_deref(),
+                &artist,
+                title.as_deref(),
+                tagged_mbid,
+            );
+            let _ = tx.send(UpdateMessage::AlbumArt {
+                cache_key,
+                cover_url,
+            });
+        });
+
+        None
+    }
+
+    fn cache_missing_album_art(&mut self, song: &mpd::Song) -> Option<String> {
+        let cache_key = format!("file:{}", song.file);
+        if let Some(cached) = self.album_art_cache.get(&cache_key) {
+            return cached.clone();
+        }
+
+        self.album_art_cache.insert(cache_key, None);
+        None
+    }
+
+    fn resolve_coverart_url(url: String) -> String {
+        let client = match HttpClient::builder().build() {
+            Ok(client) => client,
+            Err(_) => return url,
+        };
+
+        if let Ok(response) = client.head(&url).send() {
+            return response.url().to_string();
+        }
+
+        if let Ok(response) = client.get(&url).send() {
+            return response.url().to_string();
+        }
+
+        url
+    }
+
+    fn fallback_variants(value: &str) -> Vec<String> {
+        let mut variants = Vec::new();
+        let mut seen = HashSet::new();
+
+        let trimmed = value.trim();
+        if seen.insert(trimmed.to_string()) {
+            variants.push(trimmed.to_string());
+        }
+
+        if let Some(without_parens) = Self::strip_parenthetical_suffix(trimmed) {
+            if seen.insert(without_parens.clone()) {
+                variants.push(without_parens);
+            }
+        }
+
+        if let Some(without_ep) = Self::strip_dash_suffix(trimmed, "ep") {
+            if seen.insert(without_ep.clone()) {
+                variants.push(without_ep);
+            }
+        }
+
+        variants
+    }
+
+    fn strip_parenthetical_suffix(value: &str) -> Option<String> {
+        let trimmed = value.trim();
+        if trimmed.ends_with(')') {
+            if let Some(idx) = trimmed.rfind(" (") {
+                let without = trimmed[..idx].trim();
+                if !without.is_empty() && without != trimmed {
+                    return Some(without.to_string());
+                }
+            }
+        }
+        None
+    }
+
+    fn strip_dash_suffix(value: &str, suffix: &str) -> Option<String> {
+        let trimmed = value.trim();
+        let suffix_lower = format!(" - {}", suffix.to_ascii_lowercase());
+        let trimmed_lower = trimmed.to_ascii_lowercase();
+        if trimmed_lower.ends_with(&suffix_lower) {
+            let without = trimmed[..trimmed.len() - suffix_lower.len()].trim();
+            if !without.is_empty() && without != trimmed {
+                return Some(without.to_string());
+            }
+        }
+        None
+    }
+
+    fn fetch_album_art_blocking(
+        album: Option<&str>,
+        artist: &str,
+        title: Option<&str>,
+        tagged_mbid: Option<String>,
+    ) -> Option<String> {
+        let mbid = match tagged_mbid {
+            Some(mbid) => Some(mbid),
+            None => match album {
+                Some(album) => {
+                    let mut found = None;
+                    for candidate in Self::fallback_variants(album) {
+                        let query = ReleaseSearchQuery::query_builder()
+                            .artist(artist)
+                            .and()
+                            .release(&candidate)
+                            .build();
+                        let results = musicbrainz_rs::entity::release::Release::search(query)
+                            .execute()
+                            .ok();
+                        found = results
+                            .and_then(|results| results.entities.into_iter().next())
+                            .map(|release| release.id);
+                        if found.is_some() {
+                            break;
+                        }
+                    }
+                    found
+                }
+                None => title.and_then(|title| {
+                    let mut found = None;
+                    for candidate in Self::fallback_variants(title) {
+                        let query = RecordingSearchQuery::query_builder()
+                            .artist(artist)
+                            .and()
+                            .recording(&candidate)
+                            .build();
+                        let results = Recording::search(query).with_releases().execute().ok();
+                        found = results
+                            .and_then(|results| results.entities.into_iter().next())
+                            .and_then(|recording| recording.releases)
+                            .and_then(|mut releases| releases.into_iter().next())
+                            .map(|release| release.id);
+                        if found.is_some() {
+                            break;
+                        }
+                    }
+                    found
+                }),
+            },
+        };
+
+        let Some(mbid) = mbid else {
+            return None;
+        };
+
+        let cover = match musicbrainz_rs::entity::release::Release::fetch_coverart()
+            .id(&mbid)
+            .execute()
+        {
+            Ok(cover) => cover,
+            Err(_) => return None,
+        };
+
+        let cover_url = match cover {
+            musicbrainz_rs::entity::CoverartResponse::Url(url) => Some(url),
+            musicbrainz_rs::entity::CoverartResponse::Json(json) => json
+                .images
+                .iter()
+                .find(|image| image.front)
+                .or_else(|| json.images.first())
+                .map(|image| image.image.clone()),
+        };
+
+        cover_url.map(Self::resolve_coverart_url)
+    }
+
+    fn run_update(&mut self) {
+        let status = self.mpdc.lock().unwrap().status().unwrap();
+        let song = self.mpdc.lock().unwrap().currentsong().unwrap();
 
         let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
 
@@ -87,15 +347,10 @@ impl Application {
         let started_ms = start_dur.as_millis() as u64;
         let ends_ms = end_dur.as_millis() as u64;
 
-        let song = self.mpdc.currentsong().unwrap();
         if let Some(song) = song {
-            //let cover = self.mpdc.albumart(&song);
-            //let mut bytes = Vec::<u8>::new();
-            //if let Ok(cover) = cover {
-            //    bytes = cover;
-            //}
             if status.state == State::Play {
                 let album = Self::tag(&song, "Album");
+                let album_art = self.fetch_album_art(&song);
 
                 let _err = self.client.set_activity(|a| {
                     a.state(format!(
@@ -106,6 +361,23 @@ impl Application {
                     .details(song.title.unwrap_or("Unknown".into()))
                     .activity_type(ActivityType::Listening)
                     .timestamps(|_| ActivityTimestamps::new().start(started_ms).end(ends_ms))
+                    .assets(|_| {
+                        let mut assets = ActivityAssets::new()
+                            .small_image(self.config.small_image.clone().unwrap_or(MPD_LOGO.into()))
+                            .small_text(
+                                self.config
+                                    .small_text
+                                    .clone()
+                                    .unwrap_or("Music Player Daemon".into()),
+                            );
+                        if let Some(album_art) = album_art {
+                            eprintln!("URL: {}", album_art);
+                            assets = assets.large_image(album_art);
+                        } else {
+                            eprintln!("No album art URL found");
+                        }
+                        assets
+                    })
                 });
             } else {
                 let _err = self.client.clear_activity();
@@ -113,7 +385,75 @@ impl Application {
         } else {
             let _err = self.client.clear_activity();
         }
-        sleep(Duration::from_secs(1));
+    }
+
+    fn run(&mut self) {
+        self.running.store(true, Ordering::Relaxed);
+
+        let running = self.running.clone();
+        let mpdc = self.mpdc.clone();
+
+        {
+            let tx = self.update_tx.clone();
+            let running = running.clone();
+            std::thread::spawn(move || {
+                let mut last_state: Option<State> = None;
+                let mut last_song_id: Option<Id> = None;
+                while running.load(Ordering::Relaxed) {
+                    let _ = mpdc.lock().unwrap().idle(&[Subsystem::Player]);
+                    if !running.load(Ordering::Relaxed) {
+                        break;
+                    }
+
+                    let status = mpdc.lock().unwrap().status().unwrap();
+                    let song_id = status.song.map(|place| place.id);
+
+                    if last_state != Some(status.state) || last_song_id != song_id {
+                        last_state = Some(status.state);
+                        last_song_id = song_id;
+                        let _ = tx.send(UpdateMessage::Player);
+                    }
+                }
+            });
+        }
+
+        {
+            let tx = self.update_tx.clone();
+            let running = running.clone();
+            ctrlc::set_handler(move || {
+                running.store(false, Ordering::Relaxed);
+                let _ = tx.send(UpdateMessage::Stop);
+            })
+            .expect("Failed to set Ctrl-C handler");
+        }
+
+        self.run_update();
+
+        while self.running.load(Ordering::Relaxed) {
+            let Ok(message) = self.update_rx.recv() else {
+                break;
+            };
+            match message {
+                UpdateMessage::Player => self.run_update(),
+                UpdateMessage::AlbumArt {
+                    cache_key,
+                    cover_url,
+                } => {
+                    self.album_art_pending.remove(&cache_key);
+                    self.album_art_cache.insert(cache_key, cover_url);
+                    self.run_update();
+                }
+                UpdateMessage::Stop => break,
+            }
+        }
+
+        self.run_update();
+    }
+}
+
+impl Drop for Application {
+    fn drop(&mut self) {
+        let _ = self.client.clear_activity();
     }
 }
 
@@ -126,8 +466,5 @@ fn main() {
         .unwrap_or_default();
 
     let mut app = Application::new(config);
-
-    loop {
-        app.run_update();
-    }
+    app.run();
 }