eva-pwm-cohost/src/artifacts/musicbrainz.rs

use musicbrainz_rs::entity::artist_credit::ArtistCredit;
use musicbrainz_rs::entity::release::Release;
use musicbrainz_rs::{ApiEndpointError, entity::recording::Recording};
use musicbrainz_rs::prelude::*;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::artifacts::tools::{DataSource, ToolDescription};
use crate::artifacts::{Album, Artifact, ArtifactBuilder, Artist, Contents, Merge, SourceID, Track};

impl From<Recording> for Track {
    fn from(value: Recording) -> Self {
        let artist = if let Some(artist) = value.artist_credit.unwrap_or_default().first() {
            Some(artist.name.clone())
        } else {
            None
        };

        let album = if let Some(album) = value.releases.unwrap_or_default().first() {
            Some(album.title.clone())
        } else {
            None
        };

        Self {
            title: value.title,
            artist,
            album,
            ..Default::default()
        }
    }
}

impl From<Release> for Album {
    fn from(value: Release) -> Self {
        let artist = if let Some(artist) = value.artist_credit.unwrap_or_default().first() {
            Some(artist.name.clone())
        } else {
            None
        }.unwrap_or_default();

        Self {
            about: value.annotation,
            title: value.title,
            artist,
            ..Default::default()
        }
    }
}

impl From<ArtistCredit> for Artist {
    fn from(value: ArtistCredit) -> Self {
        Self {
            bio: value.artist.annotation,
            location: value.artist.country,
            name: value.name,
            ..Default::default()
        }
    }
}

pub struct MBQuery;

impl MBQuery {
    fn extract_recording_data(track: Recording) -> (Artifact, Vec<Artifact>) {
        let mut ret = vec![];

        let ret_track = ArtifactBuilder::new(SourceID::Musicbrainz)
            .contents(Track::from(track.clone()))
            .mbid(Uuid::parse_str(&track.id).unwrap())
            .build();

        for release in track.releases.unwrap_or_default() {
            log::debug!("Found new release: {:?}", release);
            ret.push(ArtifactBuilder::new(SourceID::Musicbrainz)
                .mbid(Uuid::parse_str(&release.id).unwrap())
                .contents(Album::from(release))
                .build());
        }
        for artist in track.artist_credit.unwrap_or_default() {
            ret.push(ArtifactBuilder::new(SourceID::Musicbrainz)
                .mbid(Uuid::parse_str(&artist.artist.id).unwrap())
                .contents(Artist::from(artist))
                .build());
        }

        (ret_track, ret)
    }
}

impl DataSource for MBQuery {
    type Error = ApiEndpointError;
    type Args = MusicbrainzQueryArgs;

    async fn synchronize(&mut self, artifact: &mut Artifact) -> Result<Vec<Artifact>, Self::Error> {
        let mut ret = vec![];
        if artifact.mbid.is_none() {
            return Ok(ret);
        }
        let artifact_id = artifact.mbid.clone().unwrap();
        log::debug!("Synchronizing {} with musicbrainz", artifact_id);
        match artifact.contents {
            Contents::Track(ref mut target_track) => {
                let mb_track = Recording::fetch()
                    .id(&artifact_id.to_string())
                    .with_releases().with_artists().with_annotations().execute_async().await;

                let track = match mb_track {
                    Ok(track) => track,
                    Err(err) => {
                        log::error!("Failed to grab musicbrainz data: {:?}", err);
                        return Err(err);
                    }
                };

                let (track, mut new_artifacts) = Self::extract_recording_data(track);

                ret.push(track.clone());
                ret.append(&mut new_artifacts);

                artifact.sources.insert(SourceID::Musicbrainz);

                if let Contents::Track(track) = track.contents {
                        target_track.merge(track);
                }
            },
            _ => ()
        }

        Ok(ret)
    }

    async fn query(&mut self, args: &Self::Args) -> Result<Vec<Artifact>, Self::Error> {
        let mut ret = vec![];

        for mbid in &args.mb_ids {
            log::debug!("Fetching recording id {}", mbid);
            let track = Recording::fetch()
                .id(&mbid)
                .with_releases().with_artists().with_annotations().execute_async().await;

            let track = match track {
                Ok(track) => track,
                Err(err) => {
                    log::error!("Failed to grab musicbrainz data: {:?}", err);
                    continue;
                }
            };

            let (track, mut new_artifacts) = Self::extract_recording_data(track);

            ret.push(track);
            ret.append(&mut new_artifacts);
        }

        Ok(ret)
    }
}

impl ToolDescription for MBQuery {
    fn description(&self) -> &str {
        "Fetches artifacts from Musicbrainz"
    }
    
    fn name(&self) -> &str {
        "query_musicbrainz"
    }
}

#[derive(Debug, Default, Deserialize, Serialize, JsonSchema)]
pub struct MusicbrainzQueryArgs {
    pub mb_ids: Vec<String>
}