path: root/import/src/plugins/musicbrainz.rs

/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/

use crate::{
    USER_AGENT,
    helpers::get_or_insert_slug,
    plugins::{
        ImportPlugin, PluginContext, PluginInfo,
        musicbrainz::reltypes::{VGMDB, WIKIDATA},
    },
    source_rank::ObjectImportSourceExt,
};
use anyhow::{Context, Result};
use jellycache::Cache;
use jellycommon::{jellyobject::ObjectBuffer, *};
use jellydb::RowNum;
use log::info;
use reqwest::{
    Client, ClientBuilder,
    header::{HeaderMap, HeaderName, HeaderValue},
};
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use tokio::{
    runtime::Handle,
    sync::Semaphore,
    time::{Instant, sleep_until},
};

pub mod reltypes {
    pub const MUSIC_VIDEO: &str = "ce3de655-7451-44d1-9224-87eb948c205d";
    pub const INSTRUMENTAL: &str = "9fc01a58-7801-4bd2-b07d-61cc7ffacf90";
    pub const VOCAL: &str = "0fdbe3c6-7700-4a31-ae54-b53f06ae1cfa";
    pub const RECORDING: &str = "a01ee869-80a8-45ef-9447-c59e91aa7926";
    pub const PROGRAMMING: &str = "36c50022-44e0-488d-994b-33f11d20301e";
    pub const PRODUCER: &str = "5c0ceac3-feb4-41f0-868d-dc06f6e27fc0";
    pub const ARTIST: &str = "5c0ceac3-feb4-41f0-868d-dc06f6e27fc0";
    pub const PHONOGRAPHIC_COPYRIGHT: &str = "7fd5fbc0-fbf4-4d04-be23-417d50a4dc30";
    pub const MIX: &str = "3e3102e1-1896-4f50-b5b2-dd9824e46efe";
    pub const INSTRUMENT: &str = "59054b12-01ac-43ee-a618-285fd397e461";
    pub const WIKIDATA: &str = "689870a4-a1e4-4912-b17f-7b2664215698";
    pub const VGMDB: &str = "0af15ab3-c615-46d6-b95b-a5fcd2a92ed9";
}

pub struct MusicBrainz {
    client: Client,
    rate_limit: Arc<Semaphore>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbRecordingRel {
    pub id: String,
    pub first_release_date: Option<String>,
    pub title: String,
    pub isrcs: Vec<String>,
    pub video: bool,
    pub disambiguation: String,
    pub length: Option<u32>,
    pub relations: Vec<MbRelation>,
    pub artist_credit: Vec<MbArtistCredit>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbArtistRel {
    pub id: String,
    pub isnis: Vec<String>,
    pub ipis: Vec<String>,
    pub name: String,
    pub disambiguation: String,
    pub country: Option<String>,
    pub sort_name: String,
    pub gender_id: Option<String>,
    pub area: Option<MbArea>,
    pub begin_area: Option<MbArea>,
    pub end_area: Option<MbArea>,
    pub life_span: MbTimespan,
    pub relations: Vec<MbRelation>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbArtistCredit {
    pub name: String,
    pub artist: MbArtist,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbRelation {
    pub direction: String,
    pub r#type: String,
    pub type_id: String,
    pub begin: Option<String>,
    pub end: Option<String>,
    pub ended: bool,
    pub target_type: String,
    pub target_credit: String,
    pub source_credit: String,
    pub attributes: Vec<String>,
    pub attribute_ids: BTreeMap<String, String>,
    pub attribute_values: BTreeMap<String, String>,

    pub work: Option<MbWork>,
    pub artist: Option<MbArtist>,
    pub url: Option<MbUrl>,
    pub recording: Option<MbRecording>,
    pub series: Option<MbSeries>,
    pub event: Option<MbEvent>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbSeries {
    pub id: String,
    pub r#type: Option<String>,
    pub type_id: Option<String>,
    pub name: String,
    pub disambiguation: String,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbRecording {
    pub id: String,
    pub title: String,
    #[serde(default)]
    pub isrcs: Vec<String>,
    pub video: bool,
    pub disambiguation: String,
    pub length: Option<u32>,
    #[serde(default)]
    pub artist_credit: Vec<MbArtistCredit>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbWork {
    pub id: String,
    pub r#type: Option<String>,
    pub type_id: Option<String>,
    pub languages: Vec<String>,
    pub iswcs: Vec<String>,
    pub language: Option<String>,
    pub title: String,
    pub attributes: Vec<String>,
    pub disambiguation: String,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbEvent {
    pub id: String,
    pub r#type: Option<String>,
    pub type_id: Option<String>,
    pub name: String,
    pub time: String,
    pub cancelled: bool,
    pub setlist: String,
    pub life_span: MbTimespan,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbArtist {
    pub id: String,
    pub r#type: Option<String>,
    pub type_id: Option<String>,
    pub name: String,
    pub disambiguation: String,
    pub country: Option<String>,
    pub sort_name: String,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbTimespan {
    pub begin: Option<String>,
    pub end: Option<String>,
    pub ended: bool,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbArea {
    pub name: String,
    pub sort_name: String,
    #[serde(default)]
    pub iso_3166_1_codes: Vec<String>,
    pub id: String,
    pub disambiguation: String,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct MbUrl {
    pub id: String,
    pub resource: String,
}

impl Default for MusicBrainz {
    fn default() -> Self {
        Self::new()
    }
}

impl MusicBrainz {
    const MAX_PAR_REQ: usize = 4;
    pub fn new() -> Self {
        let client = ClientBuilder::new()
            .default_headers(HeaderMap::from_iter([
                (
                    HeaderName::from_static("accept"),
                    HeaderValue::from_static("application/json"),
                ),
                (
                    HeaderName::from_static("user-agent"),
                    HeaderValue::from_static(USER_AGENT),
                ),
            ]))
            .build()
            .unwrap();
        Self {
            client,
            // send at most 1 req/s according to musicbrainz docs, each lock is held for 10s
            // this implementation also never sends more than MAX_PAR_REQ requests in-flight.
            rate_limit: Arc::new(Semaphore::new(Self::MAX_PAR_REQ)),
        }
    }

    pub fn lookup_recording(
        &self,
        cache: &Cache,
        id: &str,
        rt: &Handle,
    ) -> Result<Arc<MbRecordingRel>> {
        cache
            .cache_memory(&format!("ext/musicbrainz/recording/{id}.json"), move || {
                rt.block_on(async {
                    let _permit = self.rate_limit.clone().acquire_owned().await?;
                    let permit_drop_ts =
                        Instant::now() + Duration::from_secs(Self::MAX_PAR_REQ as u64);
                    info!("recording lookup: {id}");

                    let inc = [
                        "isrcs",
                        "artists",
                        "area-rels",
                        "artist-rels",
                        "event-rels",
                        "genre-rels",
                        "instrument-rels",
                        "label-rels",
                        "place-rels",
                        "recording-rels",
                        "release-rels",
                        "release-group-rels",
                        "series-rels",
                        "url-rels",
                        "work-rels",
                    ]
                    .join("+");

                    let resp = self
                        .client
                        .get(format!(
                            "https://musicbrainz.org/ws/2/recording/{id}?inc={inc}"
                        ))
                        .send()
                        .await?
                        .error_for_status()?
                        .json::<MbRecordingRel>()
                        .await?;

                    tokio::task::spawn(async move {
                        sleep_until(permit_drop_ts).await;
                        drop(_permit);
                    });

                    Ok(resp)
                })
            })
            .context("musicbrainz recording lookup")
    }

    pub fn lookup_artist(&self, cache: &Cache, id: &str, rt: &Handle) -> Result<Arc<MbArtistRel>> {
        cache
            .cache_memory(&format!("ext/musicbrainz/artist/{id}.json"), move || {
                rt.block_on(async {
                    let _permit = self.rate_limit.clone().acquire_owned().await?;
                    let permit_drop_ts =
                        Instant::now() + Duration::from_secs(Self::MAX_PAR_REQ as u64);
                    info!("artist lookup: {id}");

                    let inc = [
                        "area-rels",
                        "artist-rels",
                        "event-rels",
                        "genre-rels",
                        "instrument-rels",
                        "label-rels",
                        "place-rels",
                        "recording-rels",
                        "release-rels",
                        "release-group-rels",
                        "series-rels",
                        "url-rels",
                        "work-rels",
                    ]
                    .join("+");

                    let resp = self
                        .client
                        .get(format!(
                            "https://musicbrainz.org/ws/2/artist/{id}?inc={inc}"
                        ))
                        .send()
                        .await?
                        .error_for_status()?
                        .json::<MbArtistRel>()
                        .await?;

                    tokio::task::spawn(async move {
                        sleep_until(permit_drop_ts).await;
                        drop(_permit);
                    });

                    Ok(resp)
                })
            })
            .context("musicbrainz artist lookup")
    }
}

impl ImportPlugin for MusicBrainz {
    fn info(&self) -> PluginInfo {
        PluginInfo {
            name: "musicbrainz",
            tag: MSOURCE_MUSICBRAINZ,
            handle_process: true,
            ..Default::default()
        }
    }
    fn process(&self, ct: &PluginContext, node_row: RowNum) -> Result<()> {
        self.process_recording(ct, node_row)?;
        self.process_artist(ct, node_row)?;
        Ok(())
    }
}
impl MusicBrainz {
    fn process_recording(&self, ct: &PluginContext, node_row: RowNum) -> Result<()> {
        let data = ct.ic.get_node(node_row)?.unwrap();
        let data = data.as_object();

        let Some(mbid) = data
            .get(NO_IDENTIFIERS)
            .unwrap_or_default()
            .get(IDENT_MUSICBRAINZ_RECORDING)
        else {
            return Ok(());
        };

        let rec = self.lookup_recording(&ct.ic.cache, mbid, ct.rt)?;

        ct.ic.db.transaction(&mut |txn| {
            let mut node = txn.get(node_row)?.unwrap();

            node = node.as_object().insert_s(ct.is, NO_TITLE, &rec.title);
            if let Some(a) = rec.artist_credit.first() {
                node = node
                    .as_object()
                    .insert_s(ct.is, NO_SUBTITLE, &a.artist.name);
            }
            node = node.as_object().update(NO_IDENTIFIERS, |ids| {
                ids.insert_multi(
                    IDENT_ISRC,
                    &rec.isrcs.iter().map(|e| e.as_str()).collect::<Vec<_>>(),
                )
            });

            for rel in &rec.relations {
                use reltypes::*;
                let Some((role, cat)) = (match rel.type_id.as_str() {
                    INSTRUMENT => Some(("", CRCAT_INSTRUMENT)),
                    VOCAL => Some(("", CRCAT_VOCAL)),
                    PRODUCER => Some(("", CRCAT_PRODUCER)),
                    MIX => Some(("mix ", CRCAT_ENGINEER)),
                    PHONOGRAPHIC_COPYRIGHT => Some(("phonographic copyright ", CRCAT_ENGINEER)),
                    PROGRAMMING => Some(("programming ", CRCAT_ENGINEER)),
                    _ => None,
                }) else {
                    continue;
                };

                let artist = rel.artist.as_ref().unwrap();
                let artist_slug = format!("musicbrainz-artist-{}", artist.id);
                let artist_row = get_or_insert_slug(txn, &artist_slug)?;

                let mut artist_node = txn.get(artist_row)?.unwrap();
                artist_node = artist_node.as_object().update(NO_IDENTIFIERS, |ids| {
                    ids.insert(IDENT_MUSICBRAINZ_ARTIST, &artist.id)
                });
                artist_node = artist_node.as_object().insert(NO_KIND, KIND_PERSON);
                txn.update(artist_row, artist_node)?;
                ct.pending_nodes.lock().unwrap().insert(artist_row);

                let credit = ObjectBuffer::new(&mut [
                    (CR_NODE.0, &artist_row),
                    (CR_KIND.0, &cat),
                    (CR_ROLE.0, &role),
                ]);
                node = node
                    .as_object()
                    .extend_object(NO_CREDIT, CR_NODE.0, Some(credit))
            }

            txn.update(node_row, node)
        })?;

        Ok(())
    }

    fn process_artist(&self, ct: &PluginContext, node_row: RowNum) -> Result<()> {
        let data = ct.ic.get_node(node_row)?.unwrap();
        let data = data.as_object();

        let Some(mbid) = data
            .get(NO_IDENTIFIERS)
            .unwrap_or_default()
            .get(IDENT_MUSICBRAINZ_ARTIST)
        else {
            return Ok(());
        };

        let artist = self.lookup_artist(&ct.ic.cache, mbid, ct.rt)?;

        ct.ic.db.transaction(&mut |txn| {
            let mut node = txn.get(node_row)?.unwrap();
            node = node.as_object().insert_s(ct.is, NO_TITLE, &artist.name);

            for rel in &artist.relations {
                let url = rel.url.as_ref().map(|u| u.resource.clone());
                match rel.type_id.as_str() {
                    WIKIDATA => {
                        if let Some(url) = url
                            && let Some(id) = url.strip_prefix("https://www.wikidata.org/wiki/")
                        {
                            node = node
                                .as_object()
                                .update(NO_IDENTIFIERS, |ids| ids.insert(IDENT_WIKIDATA, id))
                        }
                    }
                    VGMDB => {
                        if let Some(url) = url
                            && let Some(id) = url.strip_prefix("https://vgmdb.net/artist/")
                        {
                            node = node
                                .as_object()
                                .update(NO_IDENTIFIERS, |ids| ids.insert(IDENT_VGMDB_ARTIST, id))
                        }
                    }
                    _ => (),
                }
            }

            txn.update(node_row, node)
        })?;

        Ok(())
    }
}