immich-tools/src/commands/auto_create_albums.rs

use std::path::{Path, PathBuf};

use crate::{
    actions::{
        action::Action,
        add_assets_to_album::{AddAssetsToAlbum, AddAssetsToAlbumArgs},
        create_album::{CreateAlbum, CreateAlbumArgs},
        fetch_all_albums::FetchAllAlbums,
        fetch_all_libraries::FetchAllLibraries,
        fetch_library_assets::FetchLibraryAssets,
    },
    context::Context,
    models::{asset::Asset, library::Library},
};
use color_eyre::eyre::Result;
use log::*;
use multimap::MultiMap;

pub async fn auto_create_albums(ctx: Context, separator: String) -> Result<()> {
    let albums = FetchAllAlbums::new(()).execute(&ctx).await?;
    let libraries = FetchAllLibraries::new(()).execute(&ctx).await?;
    let mut assets: Vec<Asset> = vec![];

    for library in &libraries {
        assets.extend(
            FetchLibraryAssets::new(library.clone())
                .execute(&ctx)
                .await?
                .drain(0..),
        );
    }

    let mut sorted_assets = MultiMap::new();

    for asset in &assets {
        let path = extract_path(asset, &libraries);
        if let Some(path) = path {
            let asset_albums = extract_album_names(path, separator.clone());
            for album in asset_albums {
                sorted_assets.insert(album, asset.clone());
            }
        }
    }

    debug!("Generated albums: {:?}", sorted_assets.keys());

    let missing_albums: Vec<_> = sorted_assets
        .keys()
        .filter(|k| !albums.iter().any(|album| album.name == **k))
        .cloned()
        .collect();

    info!("Creating missing albums: {missing_albums:?}");

    for missing_album in &missing_albums {
        let assets = sorted_assets
            .remove(&missing_album.to_string())
            .unwrap_or(vec![]);
        CreateAlbum::new(CreateAlbumArgs {
            name: missing_album.to_string(),
            assets,
        })
        .execute(&ctx)
        .await?;
    }

    for (album, assets) in sorted_assets {
        let existing_album = albums
            .iter()
            .find(|existing_album| existing_album.name == album);

        // Albums not present in albums were created just above with all the assets, so it's fine
        // to skip. On the other hand, Immich won't re-add the same asset multiple times to the
        // same album, so it's safe to re-add.
        if let Some(existing_album) = existing_album {
            info!("Adding {} assets to {}", assets.len(), album);
            AddAssetsToAlbum::new(AddAssetsToAlbumArgs {
                album: existing_album.clone(),
                assets,
            })
            .execute(&ctx)
            .await?;
        }
    }

    Ok(())
}

fn extract_path(asset: &Asset, libraries: &[Library]) -> Option<PathBuf> {
    for library in libraries {
        for import_path in &library.import_paths {
            if asset.original_path.starts_with(import_path) {
                return asset
                    .original_path
                    .strip_prefix(import_path)
                    .ok()
                    .and_then(Path::parent)
                    .map(Path::to_path_buf);
            }
        }
    }

    None
}

fn extract_album_names(folder_path: PathBuf, separator: String) -> Vec<String> {
    let mut components = vec![];
    let mut current_component = String::new();

    for component in folder_path.components() {
        let component_str = component.as_os_str().to_str().unwrap();
        current_component += component_str;
        components.push(current_component.clone());
        current_component += &separator;
    }

    components
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn extract_album_names_root() {
        let folder_path = PathBuf::new();
        let separator = String::from("//");

        let names = extract_album_names(folder_path, separator);

        assert_eq!(names, vec![] as Vec<String>);
    }

    #[test]
    fn extract_album_names_first_level() {
        let folder_path = PathBuf::from("My holiday photos");
        let separator = String::from("//");

        let names = extract_album_names(folder_path, separator);

        assert_eq!(names, vec!["My holiday photos"]);
    }

    #[test]
    fn extract_album_names_third_level() {
        let folder_path = PathBuf::from("My holiday photos/Europe/Toulouse");
        let separator = String::from("//");

        let names = extract_album_names(folder_path, separator);

        assert_eq!(
            names,
            vec![
                "My holiday photos",
                "My holiday photos//Europe",
                "My holiday photos//Europe//Toulouse",
            ]
        );
    }
}