immich-app/src/api.rs

use std::{
    collections::HashMap,
    io::Cursor,
    iter::repeat,
    sync::{Arc, Mutex},
};

use anyhow::{Context as _, anyhow};
use image::{
    DynamicImage, EncodableLayout,
    codecs::webp::{WebPDecoder, WebPEncoder},
};
use immich_sdk::{AssetId, AssetVisibility};
use slint::{Rgba8Pixel, SharedPixelBuffer};

use crate::{
    cachemap::{AsyncFnFetcher, CacheMap, Fetcher},
    thumbhash::thumbhashes_to_pixels,
};

pub type TimeBucketKey = String;

pub struct Api {
    client: immich_sdk::Client,
    buckets: Mutex<HashMap<TimeBucketKey, Arc<TimeBucket>>>,
    thumbnails: Mutex<CacheMap<AssetId, Arc<AssetThumbnail>>>,
}

fn serialize_thumbnail(_id: &AssetId, thumbnail: &Arc<AssetThumbnail>) -> Vec<u8> {
    let image: &SharedPixelBuffer<Rgba8Pixel> = &thumbnail.thumbnail;

    let mut webp = vec![];
    WebPEncoder::new_lossless(&mut webp)
        .encode(
            image.as_bytes(),
            image.width(),
            image.height(),
            image::ExtendedColorType::Rgba8,
        )
        .expect("width and height matches image.as_bytes().len()");

    webp
}
fn deserialize_thumbnail(&id: &AssetId, bytes: &[u8]) -> anyhow::Result<Arc<AssetThumbnail>> {
    let image = WebPDecoder::new(Cursor::new(bytes))
        .and_then(DynamicImage::from_decoder)
        .context("Failed to decode image")?
        .to_rgba8();

    let image = SharedPixelBuffer::<Rgba8Pixel>::clone_from_slice(
        image.as_bytes(),
        image.width(),
        image.height(),
    );

    Ok(Arc::new(AssetThumbnail {
        id,
        thumbnail: image,
    }))
}

impl Api {
    pub fn new(client: immich_sdk::Client) -> Arc<Self> {
        let client_ = client.clone();
        let thumbnail_fetcher = AsyncFnFetcher::new(move |asset_id: &AssetId| {
            let client = client_.clone();
            let asset_id = asset_id.clone();
            tokio::spawn(async move {
                let response = client
                    .assets()
                    .thumbnail(asset_id)
                    .size(immich_sdk::AssetMediaSize::Thumbnail)
                    .execute()
                    .await
                    .context(anyhow!("Failed to get asset thumbnail for {asset_id}"))?;

                let thumbnail = response
                    .decode()
                    .context(anyhow!("Failed to decode asset thumbnail for {asset_id}"))?
                    .into_rgba8();

                let pixel_buffer = SharedPixelBuffer::<Rgba8Pixel>::clone_from_slice(
                    &thumbnail,
                    thumbnail.width(),
                    thumbnail.height(),
                );

                Ok(Arc::new(AssetThumbnail {
                    id: asset_id,
                    thumbnail: pixel_buffer,
                }))
            })
        });
        let cache_map = CacheMap::new(
            Arc::new(thumbnail_fetcher) as Arc<dyn Fetcher<Arc<AssetThumbnail>, Key = AssetId>>,
            "thumbnails",
            serialize_thumbnail,
            deserialize_thumbnail,
        )
        .unwrap();

        Arc::new(Self {
            client,
            buckets: Default::default(),
            thumbnails: Mutex::new(cache_map),
        })
    }
}

pub struct TimeBucketRef {
    pub key: TimeBucketKey,
    pub count: usize,
}

pub struct TimeBucket {
    pub key: TimeBucketKey,
    pub entries: Arc<[TimeBucketEntry]>,
}

pub struct TimeBucketEntry {
    pub id: AssetId,
    pub is_favorite: bool,
    pub is_image: bool,
    pub is_trashed: bool,
    pub ratio: f64,
    pub thumbhash: Option<SharedPixelBuffer<Rgba8Pixel>>,
    pub visibility: AssetVisibility,
}

pub struct AssetThumbnail {
    pub id: AssetId,
    pub thumbnail: SharedPixelBuffer<Rgba8Pixel>,
}

impl Api {
    pub async fn get_time_buckets(&self) -> anyhow::Result<Vec<TimeBucketRef>> {
        let buckets = self
            .client
            .timeline()
            .buckets()
            .execute()
            .await
            .context(anyhow!("Failed to fetch list of time buckets",))
            .inspect_err(|e| {
                tracing::error!("{e:?}");
            })?;

        Ok(buckets
            .into_iter()
            .map(|b| TimeBucketRef {
                key: b.time_bucket,
                count: b.count.try_into().expect("count is non-negative"),
            })
            .collect())
    }

    pub async fn get_time_bucket(
        &self,
        time_bucket: TimeBucketKey,
    ) -> anyhow::Result<Arc<TimeBucket>> {
        if let Some(time_bucket) = self.buckets.lock().unwrap().get(&time_bucket).cloned() {
            return Ok(time_bucket);
        }

        let bucket = self
            .client
            .timeline()
            .bucket(&time_bucket)
            .execute()
            .await
            .context(anyhow!("Failed to fetch time bucket {:?}", &time_bucket))
            .inspect_err(|e| {
                tracing::error!("{e:?}");
            })?;

        let thumbhashes = bucket.thumbhash.into_iter().flatten();

        let thumbhashes = thumbhashes_to_pixels(thumbhashes)
            .context("Failed to decode thumbhashes")?
            .into_iter()
            .chain(repeat(None));

        let entries = bucket
            .id
            .into_iter()
            .enumerate()
            .zip(thumbhashes)
            .map(|((i, id), thumbhash)| {
                TimeBucketEntry {
                    id,
                    is_favorite: bucket.is_favorite.get(i).copied().unwrap_or(false),
                    is_image: bucket.is_image.get(i).copied().unwrap_or(true),
                    is_trashed: bucket.is_trashed.get(i).copied().unwrap_or(false),
                    thumbhash,
                    ratio: 1.0,                                             // TODO
                    visibility: bucket.visibility.get(i).cloned().unwrap(), // TODO: no unwrap
                }
            })
            .collect();

        let bucket = Arc::new(TimeBucket {
            key: time_bucket,
            entries,
        });

        self.buckets
            .lock()
            .unwrap()
            .insert(bucket.key.clone(), bucket.clone());

        Ok(bucket)
    }

    pub async fn get_asset_thumbnail(
        &self,
        asset_id: AssetId,
    ) -> anyhow::Result<Arc<AssetThumbnail>> {
        let fetch = self.thumbnails.lock().unwrap().get(asset_id);
        fetch
            .await
            .context(anyhow!("Failed to fetch thumbnail for {asset_id:?}"))
    }
}