use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Mutex;
use std::time::Duration;
use futures_util::lock::Mutex as AsyncMutex;
use futures_util::stream::{self, StreamExt};
use crate::album_art::{AlbumArt, PlaylistState, set_album_artifact, set_playlist};
use crate::backoff::{backoff_delay, retry_after};
use crate::client::SunoClient;
use crate::clock::Clock;
use crate::error::Error;
use crate::ffmpeg::Ffmpeg;
use crate::fs::Filesystem;
use crate::http::{Http, HttpRequest};
use crate::lineage::LineageContext;
use crate::lyrics::AlignedLyrics;
use crate::manifest::{ArtifactState, Manifest, ManifestEntry};
use crate::model::Clip;
use crate::reconcile::{Action, Desired, Plan, set_manifest_artifact, set_manifest_stem};
use crate::tag::{Cover, TrackMetadata, flac_picture_data_budget, tag_flac, tag_mp3, tag_wav};
use crate::tag_alac::tag_alac;
use crate::vocab::{ArtifactKind, AudioFormat, SourceMode, StemFormat, WebpEncodeSettings};
mod artifact;
mod audio;
mod classify;
mod cover;
mod stem;
mod tag;
use classify::*;
type ClientLock<'a, C> = AsyncMutex<&'a SunoClient<C>>;
#[derive(Debug, Clone)]
pub struct ExecOptions {
pub max_retries: u32,
pub wav_poll_attempts: u32,
pub wav_poll_interval: Duration,
pub concurrency: u32,
pub embed_animated_cover: bool,
pub cover_webp: WebpEncodeSettings,
}
impl Default for ExecOptions {
fn default() -> Self {
Self {
max_retries: 3,
wav_poll_attempts: 24,
wav_poll_interval: Duration::from_secs(5),
concurrency: 4,
embed_animated_cover: false,
cover_webp: WebpEncodeSettings::default(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RunStatus {
#[default]
Completed,
AuthAborted,
DiskFull,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Failure {
pub clip_id: String,
pub reason: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ExecOutcome {
pub downloaded: usize,
pub reformatted: usize,
pub retagged: usize,
pub renamed: usize,
pub deleted: usize,
pub skipped: usize,
pub artifacts_written: usize,
pub artifacts_deleted: usize,
pub failures: Vec<Failure>,
pub status: RunStatus,
}
impl ExecOutcome {
pub fn failed(&self) -> usize {
self.failures.len()
}
fn record(&mut self, effect: Effect) {
match effect {
Effect::Downloaded => self.downloaded += 1,
Effect::Reformatted => self.reformatted += 1,
Effect::Retagged => self.retagged += 1,
Effect::Renamed => self.renamed += 1,
Effect::Deleted => self.deleted += 1,
Effect::Skipped => self.skipped += 1,
Effect::ArtifactWritten => self.artifacts_written += 1,
Effect::ArtifactDeleted => self.artifacts_deleted += 1,
}
}
}
pub struct Ports<'a, H, F, G, C> {
pub client: &'a SunoClient<C>,
pub http: &'a H,
pub fs: &'a F,
pub ffmpeg: &'a G,
pub clock: &'a C,
}
#[allow(clippy::too_many_arguments)]
pub async fn execute<H, F, G, C>(
plan: &Plan,
manifest: &mut Manifest,
albums: &mut BTreeMap<String, AlbumArt>,
playlists: &mut BTreeMap<String, PlaylistState>,
desired: &[Desired],
synced: &HashMap<String, AlignedLyrics>,
ports: Ports<'_, H, F, G, C>,
opts: &ExecOptions,
) -> ExecOutcome
where
H: Http,
F: Filesystem,
G: Ffmpeg,
C: Clock,
{
let Ports {
client,
http,
fs,
ffmpeg,
clock,
} = ports;
let by_id: HashMap<&str, &Desired> = desired.iter().map(|d| (d.clip.id.as_str(), d)).collect();
let by_path: HashMap<&str, &Desired> = desired.iter().map(|d| (d.path.as_str(), d)).collect();
let mut tracked_paths: HashMap<String, u32> = HashMap::new();
for (_, entry) in manifest.iter() {
for path in entry.artifact_paths() {
*tracked_paths.entry(path.to_owned()).or_default() += 1;
}
}
for art in albums.values() {
for state in [
art.folder_jpg.as_ref(),
art.folder_webp.as_ref(),
art.folder_mp4.as_ref(),
]
.into_iter()
.flatten()
{
*tracked_paths.entry(state.path.clone()).or_default() += 1;
}
}
for playlist in playlists.values() {
*tracked_paths.entry(playlist.path.clone()).or_default() += 1;
}
let cover_wanted: HashSet<&str> = plan
.actions
.iter()
.filter_map(|action| match action {
Action::WriteArtifact {
kind: ArtifactKind::CoverJpg,
source_url,
..
} if !source_url.is_empty() => Some(source_url.as_str()),
_ => None,
})
.collect();
let cover_cache: Mutex<HashMap<String, Vec<u8>>> = Mutex::new(HashMap::new());
let mut folder_cover_uses: HashMap<&str, u32> = HashMap::new();
for action in &plan.actions {
if let Action::WriteArtifact {
kind: ArtifactKind::FolderWebp | ArtifactKind::FolderMp4,
source_url,
..
} = action
&& !source_url.is_empty()
{
*folder_cover_uses.entry(source_url.as_str()).or_default() += 1;
}
}
let shared_cover_urls: HashSet<&str> = folder_cover_uses
.into_iter()
.filter(|(_, uses)| *uses > 1)
.map(|(url, _)| url)
.collect();
let ctx = Ctx {
http,
fs,
ffmpeg,
clock,
opts,
by_id: &by_id,
by_path: &by_path,
synced,
cover_cache: &cover_cache,
cover_wanted: &cover_wanted,
shared_cover_urls: &shared_cover_urls,
};
let mut outcome = ExecOutcome::default();
let mut committed: BTreeSet<String> = BTreeSet::new();
let client_lock = AsyncMutex::new(client);
let concurrency = opts.concurrency.max(1) as usize;
let ctx_ref = &ctx;
let client_lock_ref = &client_lock;
let pre_clip_ids: HashSet<String> = manifest.entries.keys().cloned().collect();
let audio_clip_ids: HashSet<&str> = plan
.actions
.iter()
.filter_map(|action| match action {
Action::Download { clip, .. } | Action::Reformat { clip, .. } => Some(clip.id.as_str()),
_ => None,
})
.collect();
let mut prepares = stream::iter(
plan.actions
.iter()
.filter(|action| is_prepareable(action, &pre_clip_ids, &audio_clip_ids))
.map(|action| async move { ctx_ref.prepare(client_lock_ref, action).await }),
)
.buffered(concurrency);
for action in &plan.actions {
let result = if is_prepareable(action, &pre_clip_ids, &audio_clip_ids) {
match prepares.next().await {
Some(Ok(Prepared::Audio(rendered))) => ctx.commit_audio(manifest, rendered),
Some(Ok(Prepared::Artifact(prepared))) => ctx.commit_artifact(
manifest,
albums,
playlists,
prepared,
&mut tracked_paths,
&committed,
),
Some(Ok(Prepared::Stem(prepared))) => {
ctx.commit_stem(manifest, prepared, &mut tracked_paths, &committed)
}
Some(Err(fail)) => Err(fail),
None => unreachable!("buffered yields one result per prepareable action"),
}
} else {
ctx.apply(
client_lock_ref,
action,
manifest,
albums,
playlists,
&mut tracked_paths,
&committed,
)
.await
};
match result {
Ok(effect) => {
outcome.record(effect);
if let Some(dest) = written_path(action) {
committed.insert(dest.to_owned());
}
}
Err(fail) => {
let abort = abort_status(fail.class);
outcome.failures.push(Failure {
clip_id: fail.clip_id,
reason: fail.reason,
});
if let Some(status) = abort {
outcome.status = status;
break;
}
}
}
}
drop(prepares);
let _ = fs.prune_empty_dirs("");
outcome
}
fn is_prepareable(
action: &Action,
pre_clip_ids: &HashSet<String>,
audio_clip_ids: &HashSet<&str>,
) -> bool {
match action {
Action::Download { .. } | Action::Reformat { .. } => true,
Action::WriteArtifact {
kind,
owner_id,
content: None,
..
} => {
if matches!(kind, ArtifactKind::FolderWebp | ArtifactKind::FolderMp4) {
return false;
}
if *kind == ArtifactKind::CoverJpg && audio_clip_ids.contains(owner_id.as_str()) {
return false;
}
!is_per_clip_kind(*kind) || pre_clip_ids.contains(owner_id.as_str())
}
Action::WriteStem { clip_id, .. } => pre_clip_ids.contains(clip_id.as_str()),
_ => false,
}
}
fn written_path(action: &Action) -> Option<&str> {
match action {
Action::Download { path, .. }
| Action::Reformat { path, .. }
| Action::WriteArtifact { path, .. }
| Action::WriteStem { path, .. } => Some(path),
Action::Rename { to, .. }
| Action::MoveArtifact { to, .. }
| Action::MoveStem { to, .. } => Some(to),
_ => None,
}
}
struct RenderedAudio {
clip_id: String,
path: String,
format: AudioFormat,
from_path: Option<String>,
effect: Effect,
bytes: Vec<u8>,
}
struct PreparedArtifact {
kind: ArtifactKind,
path: String,
hash: String,
owner_id: String,
bytes: Vec<u8>,
}
struct PreparedStem {
clip_id: String,
key: String,
path: String,
hash: String,
bytes: Vec<u8>,
}
enum Prepared {
Audio(RenderedAudio),
Artifact(PreparedArtifact),
Stem(PreparedStem),
}
struct EmbedCover {
bytes: Vec<u8>,
mime: &'static str,
}
impl EmbedCover {
fn as_cover(&self) -> Cover<'_> {
Cover {
bytes: &self.bytes,
mime: self.mime,
}
}
}
enum Effect {
Downloaded,
Reformatted,
Retagged,
Renamed,
Deleted,
Skipped,
ArtifactWritten,
ArtifactDeleted,
}
fn is_album_kind(kind: ArtifactKind) -> bool {
matches!(
kind,
ArtifactKind::FolderJpg | ArtifactKind::FolderWebp | ArtifactKind::FolderMp4
)
}
fn is_playlist_kind(kind: ArtifactKind) -> bool {
matches!(kind, ArtifactKind::Playlist)
}
fn is_per_clip_kind(kind: ArtifactKind) -> bool {
matches!(
kind,
ArtifactKind::CoverJpg
| ArtifactKind::CoverWebp
| ArtifactKind::DetailsTxt
| ArtifactKind::LyricsTxt
| ArtifactKind::Lrc
| ArtifactKind::VideoMp4
)
}
fn playlist_name_from_path(path: &str) -> String {
std::path::Path::new(path)
.file_stem()
.map(|stem| stem.to_string_lossy().into_owned())
.unwrap_or_default()
}
struct Ctx<'a, H, F, G, C> {
http: &'a H,
fs: &'a F,
ffmpeg: &'a G,
clock: &'a C,
opts: &'a ExecOptions,
by_id: &'a HashMap<&'a str, &'a Desired>,
by_path: &'a HashMap<&'a str, &'a Desired>,
synced: &'a HashMap<String, AlignedLyrics>,
cover_cache: &'a Mutex<HashMap<String, Vec<u8>>>,
cover_wanted: &'a HashSet<&'a str>,
shared_cover_urls: &'a HashSet<&'a str>,
}
impl<H, F, G, C> Ctx<'_, H, F, G, C>
where
H: Http,
F: Filesystem,
G: Ffmpeg,
C: Clock,
{
#[allow(clippy::too_many_arguments)]
async fn apply(
&self,
client_lock: &ClientLock<'_, C>,
action: &Action,
manifest: &mut Manifest,
albums: &mut BTreeMap<String, AlbumArt>,
playlists: &mut BTreeMap<String, PlaylistState>,
tracked_paths: &mut HashMap<String, u32>,
committed: &BTreeSet<String>,
) -> Result<Effect, Fail> {
match action {
Action::Download { .. } | Action::Reformat { .. } => {
unreachable!("audio actions are prepared concurrently")
}
Action::Retag {
clip,
lineage,
path,
} => self.retag(manifest, clip, lineage, path).await,
Action::Rename { from, to } => self.rename(manifest, from, to),
Action::Delete { path, clip_id } => self.delete(manifest, path, clip_id),
Action::Skip { clip_id } => {
self.refresh_preserve(manifest, clip_id);
Ok(Effect::Skipped)
}
Action::WriteArtifact {
kind,
path,
source_url,
hash,
owner_id,
content,
} => {
let bytes = match content.as_deref() {
Some(text) => text.as_bytes().to_vec(),
None => {
if is_per_clip_kind(*kind) && manifest.get(owner_id).is_none() {
self.cover_cache_lock().remove(source_url);
return Ok(Effect::Skipped);
}
self.artifact_bytes(*kind, source_url, owner_id).await?
}
};
self.commit_artifact(
manifest,
albums,
playlists,
PreparedArtifact {
kind: *kind,
path: path.clone(),
hash: hash.clone(),
owner_id: owner_id.clone(),
bytes,
},
tracked_paths,
committed,
)
}
Action::DeleteArtifact {
kind,
path,
owner_id,
} => self.delete_artifact(manifest, albums, playlists, *kind, path, owner_id),
Action::MoveArtifact {
kind,
from,
to,
source_url,
hash,
owner_id,
} => {
self.move_artifact(
manifest,
albums,
playlists,
*kind,
from,
to,
source_url,
hash,
owner_id,
tracked_paths,
committed,
)
.await
}
Action::WriteStem {
clip_id,
key,
stem_id,
path,
source_url,
format,
hash,
} => {
if manifest.get(clip_id).is_none() {
return Ok(Effect::Skipped);
}
let bytes = self
.fetch_stem_bytes(client_lock, clip_id, stem_id, source_url, *format)
.await?;
self.commit_stem(
manifest,
PreparedStem {
clip_id: clip_id.clone(),
key: key.clone(),
path: path.clone(),
hash: hash.clone(),
bytes,
},
tracked_paths,
committed,
)
}
Action::DeleteStem { clip_id, key, path } => {
self.delete_stem(manifest, clip_id, key, path)
}
Action::MoveStem {
clip_id,
key,
stem_id,
from,
to,
source_url,
format,
hash,
} => {
self.move_stem(
client_lock,
manifest,
clip_id,
key,
stem_id,
from,
to,
source_url,
*format,
hash,
tracked_paths,
committed,
)
.await
}
}
}
fn rename(&self, manifest: &mut Manifest, from: &str, to: &str) -> Result<Effect, Fail> {
let label = self
.by_path
.get(to)
.map(|d| d.clip.id.clone())
.unwrap_or_else(|| to.to_owned());
self.fs.rename(from, to).map_err(|err| {
if err.is_out_of_space() {
disk_fail(label, "disk full: no space left to rename")
} else {
permanent_fail(label, format!("rename failed: {err}"))
}
})?;
let clip_id = self.by_path.get(to).map(|d| d.clip.id.clone()).or_else(|| {
manifest
.entries
.iter()
.find(|(_, entry)| entry.path == from)
.map(|(id, _)| id.clone())
});
if let Some(id) = clip_id
&& let Some(entry) = manifest.entries.get_mut(&id)
{
entry.path = to.to_owned();
if let Some(d) = self.by_path.get(to) {
entry.preserve = preserve_for(d);
}
}
Ok(Effect::Renamed)
}
fn delete(&self, manifest: &mut Manifest, path: &str, clip_id: &str) -> Result<Effect, Fail> {
self.fs
.remove(path)
.map_err(|err| permanent_fail(clip_id, format!("delete failed: {err}")))?;
manifest.remove(clip_id);
Ok(Effect::Deleted)
}
async fn retry_core(&self, id: &str, err: Error, attempt: &mut u32) -> Option<Fail> {
let fail = classify_core(id, err);
if matches!(fail.class, Class::Transient) && *attempt < self.opts.max_retries {
self.clock.sleep(backoff_delay(*attempt, None)).await;
*attempt += 1;
None
} else {
Some(fail)
}
}
async fn fetch_bytes(&self, url: &str) -> Result<Vec<u8>, FetchError> {
let mut attempt: u32 = 0;
loop {
let result = self.http.send(HttpRequest::get(url)).await;
match classify_response(result) {
Ok(body) => return Ok(body),
Err(err) => {
if matches!(err.class, Class::Transient) && attempt < self.opts.max_retries {
let delay = backoff_delay(attempt, err.retry_after);
self.clock.sleep(delay).await;
attempt += 1;
continue;
}
return Err(err);
}
}
}
}
fn write_verify(&self, clip_id: &str, path: &str, bytes: &[u8]) -> Result<u64, Fail> {
self.fs.write_atomic(path, bytes).map_err(|err| {
if err.is_out_of_space() {
disk_fail(clip_id, format!("disk full: no space left to write {path}"))
} else {
permanent_fail(clip_id, format!("write failed: {err}"))
}
})?;
match self.fs.metadata(path) {
Some(stat) if stat.size == bytes.len() as u64 => Ok(stat.size),
Some(stat) => Err(permanent_fail(
clip_id,
format!("wrote {} bytes, expected {}", stat.size, bytes.len()),
)),
None => Ok(bytes.len() as u64),
}
}
fn entry(&self, clip_id: &str, path: &str, format: AudioFormat, size: u64) -> ManifestEntry {
match self.by_id.get(clip_id) {
Some(d) => manifest_entry(d, size),
None => ManifestEntry {
path: path.to_owned(),
format,
size,
..ManifestEntry::default()
},
}
}
}
fn manifest_entry(d: &Desired, size: u64) -> ManifestEntry {
ManifestEntry {
path: d.path.clone(),
format: d.format,
meta_hash: d.meta_hash.clone(),
art_hash: d.art_hash.clone(),
size,
preserve: preserve_for(d),
..Default::default()
}
}
fn preserve_for(d: &Desired) -> bool {
d.private || d.modes.contains(&SourceMode::Copy)
}
#[cfg(test)]
mod tests;