use anyhow::{Context, Result};
use chrono::{Datelike, NaiveDate};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::{Mutex, OnceCell};
use crate::api::SmugMugClient;
use crate::api::albums::{Album, ChildNode};
use crate::api::images::AlbumImage;
use crate::uploader::album_series::{AlbumSeries, AlbumSeriesBackend, MAX_ALBUM_IMAGES};
#[allow(async_fn_in_trait)]
pub trait Smug {
async fn children(&self, node_uri: &str) -> Result<Vec<ChildNode>>;
async fn create_folder(&self, parent_node_uri: &str, name: &str) -> Result<String>;
async fn create_album(&self, parent_node_uri: &str, name: &str) -> Result<Album>;
async fn image_count(&self, album: &Album) -> Result<u64>;
async fn album_images(&self, album_key: &str) -> Result<Vec<AlbumImage>>;
}
impl Smug for SmugMugClient {
async fn children(&self, node_uri: &str) -> Result<Vec<ChildNode>> {
self.list_children(node_uri).await
}
async fn create_folder(&self, parent_node_uri: &str, name: &str) -> Result<String> {
SmugMugClient::create_folder(self, parent_node_uri, name, Some("Private")).await
}
async fn create_album(&self, parent_node_uri: &str, name: &str) -> Result<Album> {
SmugMugClient::create_album(self, name, Some(parent_node_uri), "Private").await
}
async fn image_count(&self, album: &Album) -> Result<u64> {
Ok(self
.get_album(&album.album_key)
.await?
.image_count
.unwrap_or(0))
}
async fn album_images(&self, album_key: &str) -> Result<Vec<AlbumImage>> {
self.list_album_images(album_key).await
}
}
pub fn day_album_name(date: NaiveDate) -> String {
date.format("%Y-%m-%d").to_string()
}
fn month_path(date: NaiveDate) -> [String; 2] {
[
format!("{:04}", date.year()),
format!("{:02}", date.month()),
]
}
#[derive(Default)]
struct Folders {
known: HashMap<String, Option<String>>,
listings: HashMap<String, Arc<Vec<ChildNode>>>,
}
struct Inner<T> {
smug: Arc<T>,
root: Option<String>,
dry_run: bool,
folders: Mutex<Folders>,
}
impl<T: Smug> Inner<T> {
async fn listing(&self, folders: &mut Folders, node_uri: &str) -> Result<Arc<Vec<ChildNode>>> {
if let Some(listing) = folders.listings.get(node_uri) {
return Ok(listing.clone());
}
let listing = Arc::new(
self.smug
.children(node_uri)
.await
.with_context(|| format!("Failed to list folder {}", node_uri))?,
);
folders
.listings
.insert(node_uri.to_string(), listing.clone());
Ok(listing)
}
async fn folder(&self, path: &[String], create: bool) -> Result<Option<String>> {
let Some(root) = &self.root else {
return Ok(None);
};
let mut folders = self.folders.lock().await;
let mut parent = root.clone();
let mut key = String::new();
for part in path {
if !key.is_empty() {
key.push('/');
}
key.push_str(part);
let known = folders.known.get(&key).cloned();
let found = match known {
Some(Some(uri)) => Some(uri),
Some(None) if !create => return Ok(None),
_ => self
.listing(&mut folders, &parent)
.await?
.iter()
.find(|c| c.node_type == "Folder" && c.name == *part)
.map(|c| c.uri.clone()),
};
let uri = match found {
Some(uri) => uri,
None if create && !self.dry_run => {
let uri = self
.smug
.create_folder(&parent, part)
.await
.with_context(|| format!("Failed to create folder {}", key))?;
folders.listings.insert(uri.clone(), Arc::new(Vec::new()));
uri
}
None => {
folders.known.insert(key, None);
return Ok(None);
}
};
folders.known.insert(key.clone(), Some(uri.clone()));
parent = uri;
}
Ok(Some(parent))
}
async fn albums_in(&self, folder_uri: &str) -> Result<Vec<Album>> {
let mut folders = self.folders.lock().await;
Ok(self
.listing(&mut folders, folder_uri)
.await?
.iter()
.filter_map(ChildNode::album)
.collect())
}
}
pub struct DayBackend<T> {
inner: Arc<Inner<T>>,
date: NaiveDate,
}
impl<T: Smug> AlbumSeriesBackend for DayBackend<T> {
async fn find_album(&self, name: &str) -> Result<Option<Album>> {
let Some(month) = self.inner.folder(&month_path(self.date), false).await? else {
return Ok(None);
};
Ok(self
.inner
.albums_in(&month)
.await?
.into_iter()
.find(|a| a.name == name))
}
async fn image_count(&self, album: &Album) -> Result<u64> {
self.inner.smug.image_count(album).await
}
async fn create_album(&self, name: &str) -> Result<Album> {
let month = self
.inner
.folder(&month_path(self.date), true)
.await?
.context("The day's folder couldn't be created")?;
self.inner.smug.create_album(&month, name).await
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamedImage {
pub md5: Option<String>,
pub uri: String,
pub album_key: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NameCheck {
Upload(String),
Present(NamedImage),
}
pub struct Day<T> {
pub series: AlbumSeries<DayBackend<T>>,
smug: Arc<T>,
names: OnceCell<std::sync::Mutex<HashMap<String, NamedImage>>>,
}
pub fn renamed(file_name: &str, md5: &str) -> String {
let tag = &md5[..md5.len().min(8)];
match file_name.rsplit_once('.') {
Some((stem, ext)) if !stem.is_empty() => format!("{}~{}.{}", stem, tag, ext),
_ => format!("{}~{}", file_name, tag),
}
}
impl<T: Smug> Day<T> {
async fn names(&self) -> Result<&std::sync::Mutex<HashMap<String, NamedImage>>> {
self.names
.get_or_try_init(|| async {
let mut names = HashMap::new();
for existing in self.series.existing_albums().await {
let key = existing.album.album_key;
for image in
self.smug.album_images(&key).await.with_context(|| {
format!("Failed to list album {}", existing.album.name)
})?
{
names.insert(
image.file_name,
NamedImage {
md5: image.archived_md5.map(|m| m.to_lowercase()),
uri: image.uri,
album_key: key.clone(),
},
);
}
}
Ok::<_, anyhow::Error>(std::sync::Mutex::new(names))
})
.await
}
pub async fn check_name(&self, file_name: &str, md5: &str) -> Result<NameCheck> {
let names = self.names().await?;
let mut names = names.lock().unwrap();
let alternative = renamed(file_name, md5);
for candidate in [file_name, alternative.as_str()] {
match names.get(candidate) {
None => {
names.insert(
candidate.to_string(),
NamedImage {
md5: Some(md5.to_string()),
uri: String::new(),
album_key: String::new(),
},
);
return Ok(NameCheck::Upload(candidate.to_string()));
}
Some(existing)
if existing.md5.as_deref() == Some(md5) && !existing.uri.is_empty() =>
{
return Ok(NameCheck::Present(existing.clone()));
}
Some(_) => {}
}
}
let mut n = 2;
loop {
let candidate = renamed(file_name, &format!("{}-{}", md5, n));
if !names.contains_key(&candidate) {
names.insert(
candidate.clone(),
NamedImage {
md5: Some(md5.to_string()),
uri: String::new(),
album_key: String::new(),
},
);
return Ok(NameCheck::Upload(candidate));
}
n += 1;
}
}
pub async fn name_uploaded(&self, name: &str, uri: &str, album_key: &str) {
if let Some(names) = self.names.get()
&& let Some(entry) = names.lock().unwrap().get_mut(name)
{
entry.uri = uri.to_string();
entry.album_key = album_key.to_string();
}
}
pub async fn release_name(&self, name: &str) {
if let Some(names) = self.names.get() {
let mut names = names.lock().unwrap();
if names.get(name).is_some_and(|e| e.uri.is_empty()) {
names.remove(name);
}
}
}
pub async fn album_keys(&self) -> HashSet<String> {
self.series
.albums()
.await
.into_iter()
.map(|a| a.album.album_key)
.filter(|k| !k.is_empty())
.collect()
}
}
type DayCell<T> = Arc<OnceCell<Arc<Day<T>>>>;
pub struct DayAlbums<T> {
inner: Arc<Inner<T>>,
days: std::sync::Mutex<HashMap<NaiveDate, DayCell<T>>>,
}
impl<T: Smug> DayAlbums<T> {
pub fn new(smug: Arc<T>, root: Option<String>, dry_run: bool) -> Self {
DayAlbums {
inner: Arc::new(Inner {
smug,
root,
dry_run,
folders: Mutex::new(Folders::default()),
}),
days: std::sync::Mutex::new(HashMap::new()),
}
}
pub async fn day(&self, date: NaiveDate) -> Result<Arc<Day<T>>> {
let cell = self
.days
.lock()
.unwrap()
.entry(date)
.or_insert_with(|| Arc::new(OnceCell::new()))
.clone();
cell.get_or_try_init(|| async {
let backend = DayBackend {
inner: self.inner.clone(),
date,
};
let series = AlbumSeries::load(
backend,
&day_album_name(date),
MAX_ALBUM_IMAGES,
self.inner.dry_run,
)
.await
.with_context(|| format!("Failed to look up the albums for {}", date))?;
Ok::<_, anyhow::Error>(Arc::new(Day {
series,
smug: self.inner.smug.clone(),
names: OnceCell::new(),
}))
})
.await
.cloned()
}
pub fn loaded_days(&self) -> Vec<Arc<Day<T>>> {
self.days
.lock()
.unwrap()
.values()
.filter_map(|cell| cell.get().cloned())
.collect()
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use std::sync::Mutex as StdMutex;
#[derive(Default)]
pub struct FakeSmug {
pub nodes: StdMutex<HashMap<String, Vec<ChildNode>>>,
pub images: StdMutex<HashMap<String, Vec<AlbumImage>>>,
pub calls: StdMutex<Vec<String>>,
pub refuse: StdMutex<HashSet<String>>,
next: StdMutex<u32>,
}
impl FakeSmug {
pub fn id(&self) -> u32 {
let mut n = self.next.lock().unwrap();
*n += 1;
*n
}
pub fn with_root() -> (Arc<Self>, String) {
let fake = Arc::new(FakeSmug::default());
fake.nodes
.lock()
.unwrap()
.insert("/node/root".into(), Vec::new());
(fake, "/node/root".into())
}
pub fn add_folder(&self, parent: &str, name: &str) -> String {
let uri = format!("/node/f{}", self.id());
self.nodes
.lock()
.unwrap()
.get_mut(parent)
.unwrap()
.push(ChildNode {
name: name.into(),
node_type: "Folder".into(),
uri: uri.clone(),
node_id: String::new(),
url_name: String::new(),
web_uri: None,
uris: None,
});
self.nodes.lock().unwrap().insert(uri.clone(), Vec::new());
uri
}
pub fn add_album(&self, parent: &str, name: &str, images: &[(&str, &str)]) -> Album {
let id = self.id();
let key = format!("A{}", id);
let album_uri = format!("/api/v2/album/{}", key);
self.nodes
.lock()
.unwrap()
.get_mut(parent)
.unwrap()
.push(ChildNode {
name: name.into(),
node_type: "Album".into(),
uri: format!("/node/a{}", id),
node_id: format!("a{}", id),
url_name: name.into(),
web_uri: None,
uris: Some(crate::api::albums::ChildNodeUris {
album: Some(crate::api::albums::UriRef {
uri: album_uri.clone(),
}),
}),
});
let images = images
.iter()
.map(|(file, md5)| AlbumImage {
image_key: format!("{}-{}", key, file),
file_name: file.to_string(),
archived_uri: String::new(),
file_size: 1,
format: "JPG".into(),
uri: format!("{}/image/{}", album_uri, file),
title: None,
archived_md5: Some(md5.to_string()),
})
.collect();
self.images.lock().unwrap().insert(key.clone(), images);
Album {
album_key: key,
name: name.into(),
url_name: name.into(),
node_id: format!("a{}", id),
uri: album_uri,
web_uri: None,
uris: None,
image_count: None,
}
}
pub fn calls(&self, prefix: &str) -> usize {
self.calls
.lock()
.unwrap()
.iter()
.filter(|c| c.starts_with(prefix))
.count()
}
}
impl Smug for FakeSmug {
async fn children(&self, node_uri: &str) -> Result<Vec<ChildNode>> {
self.calls
.lock()
.unwrap()
.push(format!("children {}", node_uri));
self.nodes
.lock()
.unwrap()
.get(node_uri)
.cloned()
.context("no such node")
}
async fn create_folder(&self, parent: &str, name: &str) -> Result<String> {
self.calls
.lock()
.unwrap()
.push(format!("create_folder {}", name));
Ok(self.add_folder(parent, name))
}
async fn create_album(&self, parent: &str, name: &str) -> Result<Album> {
self.calls
.lock()
.unwrap()
.push(format!("create_album {}", name));
Ok(self.add_album(parent, name, &[]))
}
async fn image_count(&self, album: &Album) -> Result<u64> {
Ok(self
.images
.lock()
.unwrap()
.get(&album.album_key)
.map_or(0, |i| i.len() as u64))
}
async fn album_images(&self, album_key: &str) -> Result<Vec<AlbumImage>> {
self.calls
.lock()
.unwrap()
.push(format!("album_images {}", album_key));
Ok(self
.images
.lock()
.unwrap()
.get(album_key)
.cloned()
.unwrap_or_default())
}
}
fn d(y: i32, m: u32, day: u32) -> NaiveDate {
NaiveDate::from_ymd_opt(y, m, day).unwrap()
}
#[tokio::test]
async fn finds_existing_day_albums_without_creating_anything() {
let (fake, root) = FakeSmug::with_root();
let year = fake.add_folder(&root, "2014");
let month = fake.add_folder(&year, "07");
let album = fake.add_album(&month, "2014-07-12", &[("a.jpg", "aa")]);
let days = DayAlbums::new(fake.clone(), Some(root), false);
let day = days.day(d(2014, 7, 12)).await.unwrap();
let existing = day.series.existing_albums().await;
assert_eq!(existing.len(), 1);
assert_eq!(existing[0].album.album_key, album.album_key);
assert_eq!(fake.calls("create"), 0);
let before = fake.calls("children");
days.day(d(2014, 7, 13)).await.unwrap();
assert_eq!(fake.calls("children"), before);
}
#[tokio::test]
async fn creates_year_month_and_day_only_when_claimed() {
let (fake, root) = FakeSmug::with_root();
let days = DayAlbums::new(fake.clone(), Some(root.clone()), false);
let day = days.day(d(2020, 1, 5)).await.unwrap();
assert_eq!(fake.calls("create"), 0);
let album = day.series.claim().await.unwrap();
assert_eq!(album.name, "2020-01-05");
assert_eq!(fake.calls("create_folder"), 2);
assert_eq!(fake.calls("create_album"), 1);
let other = days.day(d(2020, 1, 6)).await.unwrap();
other.series.claim().await.unwrap();
assert_eq!(fake.calls("create_folder"), 2);
let nodes = fake.nodes.lock().unwrap();
let year = nodes[&root].iter().find(|c| c.name == "2020").unwrap();
let month = nodes[&year.uri].iter().find(|c| c.name == "01").unwrap();
let names: Vec<_> = nodes[&month.uri].iter().map(|c| c.name.clone()).collect();
assert_eq!(names, vec!["2020-01-05", "2020-01-06"]);
}
#[tokio::test]
async fn dry_run_without_root_uses_placeholders() {
let (fake, _) = FakeSmug::with_root();
let days = DayAlbums::new(fake.clone(), None, true);
let day = days.day(d(2020, 1, 5)).await.unwrap();
let album = day.series.claim().await.unwrap();
assert_eq!(album.name, "2020-01-05");
assert!(album.album_key.is_empty());
assert!(fake.calls.lock().unwrap().is_empty());
}
#[tokio::test]
async fn dry_run_creates_nothing() {
let (fake, root) = FakeSmug::with_root();
let days = DayAlbums::new(fake.clone(), Some(root), true);
let day = days.day(d(2020, 1, 5)).await.unwrap();
day.series.claim().await.unwrap();
assert_eq!(fake.calls("create"), 0);
}
#[tokio::test]
async fn names_link_same_content_and_rename_different_content() {
let (fake, root) = FakeSmug::with_root();
let year = fake.add_folder(&root, "2014");
let month = fake.add_folder(&year, "07");
fake.add_album(&month, "2014-07-12", &[("IMG_0001.JPG", "aaaa1111")]);
let days = DayAlbums::new(fake.clone(), Some(root), false);
let day = days.day(d(2014, 7, 12)).await.unwrap();
match day.check_name("IMG_0001.JPG", "aaaa1111").await.unwrap() {
NameCheck::Present(image) => assert!(image.uri.ends_with("IMG_0001.JPG")),
other => panic!("{:?}", other),
}
assert_eq!(
day.check_name("IMG_0001.JPG", "bbbb2222cccc")
.await
.unwrap(),
NameCheck::Upload("IMG_0001~bbbb2222.JPG".into())
);
assert_eq!(
day.check_name("IMG_0002.JPG", "dddd").await.unwrap(),
NameCheck::Upload("IMG_0002.JPG".into())
);
assert_eq!(
day.check_name("IMG_0002.JPG", "eeee").await.unwrap(),
NameCheck::Upload("IMG_0002~eeee.JPG".into())
);
day.release_name("IMG_0002.JPG").await;
assert_eq!(
day.check_name("IMG_0002.JPG", "ffff").await.unwrap(),
NameCheck::Upload("IMG_0002.JPG".into())
);
day.name_uploaded("IMG_0002.JPG", "/img/2", "A9").await;
assert_eq!(
day.check_name("IMG_0002.JPG", "ffff").await.unwrap(),
NameCheck::Present(NamedImage {
md5: Some("ffff".into()),
uri: "/img/2".into(),
album_key: "A9".into()
})
);
assert_eq!(fake.calls("album_images"), 1);
}
#[test]
fn renaming() {
assert_eq!(
renamed("IMG_0001.JPG", "0123456789"),
"IMG_0001~01234567.JPG"
);
assert_eq!(renamed("noext", "0123456789"), "noext~01234567");
assert_eq!(renamed(".hidden", "ab"), ".hidden~ab");
}
}