jellyfin-client/jellyfin-kiosk-rs/src/image_cache.rs

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use tokio::sync::{mpsc, RwLock, Semaphore};
use eframe::egui::{ColorImage, Context, TextureHandle, TextureOptions};

#[derive(Clone)]
pub enum ImageStatus {
    Loaded(TextureHandle),
    Loading,
    Failed,
}

fn hash_url(url: &str) -> u64 {
    use std::hash::{Hash, Hasher};
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    url.hash(&mut hasher);
    hasher.finish()
}

struct CacheInner {
    textures: HashMap<String, (u64, TextureHandle)>,
    in_flight: HashSet<String>,
    failed: HashSet<String>,
    access_counter: u64,
}

pub struct AsyncImageCache {
    inner: Mutex<CacheInner>,
    tx: mpsc::UnboundedSender<(String, Option<ColorImage>)>,
    rx: Mutex<mpsc::UnboundedReceiver<(String, Option<ColorImage>)>>,
    http_client: reqwest::Client,
    token: Arc<RwLock<Option<String>>>,
    semaphore: Arc<Semaphore>,
}

impl AsyncImageCache {
    pub fn new(token: Arc<RwLock<Option<String>>>) -> Self {
        let (tx, rx) = mpsc::unbounded_channel();
        Self {
            inner: Mutex::new(CacheInner {
                textures: HashMap::new(),
                in_flight: HashSet::new(),
                failed: HashSet::new(),
                access_counter: 0,
            }),
            tx,
            rx: Mutex::new(rx),
            http_client: reqwest::Client::new(),
            token,
            semaphore: Arc::new(Semaphore::new(10)),
        }
    }

    pub fn update(&self, ctx: &Context) {
        let mut any_new = false;
        if let Ok(mut rx_guard) = self.rx.try_lock() {
            while let Ok((url, opt_img)) = rx_guard.try_recv() {
                if let Ok(mut inner) = self.inner.lock() {
                    inner.in_flight.remove(&url);
                    match opt_img {
                        Some(color_image) => {
                            inner.access_counter += 1;
                            let counter = inner.access_counter;
                            let handle = ctx.load_texture(&url, color_image, TextureOptions::LINEAR);
                            inner.textures.insert(url, (counter, handle));
                            any_new = true;
                        }
                        None => {
                            inner.failed.insert(url);
                        }
                    }
                }
            }
        }

        // Limit cached textures to 500
        if let Ok(mut inner) = self.inner.lock() {
            if inner.textures.len() > 500 {
                let mut entries: Vec<(String, u64)> = inner
                    .textures
                    .iter()
                    .map(|(k, (c, _))| (k.clone(), *c))
                    .collect();
                entries.sort_by_key(|(_, c)| *c);
                let to_remove = entries.len() - 400;
                for (k, _) in entries.into_iter().take(to_remove) {
                    inner.textures.remove(&k);
                }
            }
        }

        if any_new {
            ctx.request_repaint();
        }
    }

    pub fn get_or_load(&self, ctx: &Context, url: &str) -> ImageStatus {
        if url.is_empty() {
            return ImageStatus::Failed;
        }

        if let Ok(mut inner) = self.inner.lock() {
            inner.access_counter += 1;
            let current_count = inner.access_counter;
            if let Some((counter, handle)) = inner.textures.get_mut(url) {
                *counter = current_count;
                return ImageStatus::Loaded(handle.clone());
            }

            if inner.failed.contains(url) {
                return ImageStatus::Failed;
            }

            if inner.in_flight.contains(url) {
                return ImageStatus::Loading;
            }

            inner.in_flight.insert(url.to_string());
        } else {
            return ImageStatus::Loading;
        }

        let url_clone = url.to_string();
        let tx = self.tx.clone();
        let client = self.http_client.clone();
        let sem = self.semaphore.clone();
        let ctx_clone = ctx.clone();
        let token_opt = self.token.try_read().ok().and_then(|g| g.clone());

        tokio::spawn(async move {
            let _permit = sem.acquire().await;
            let cache_dir = std::path::Path::new("/tmp/jellyfin-cache");
            let file_name = format!("{:016x}.cache", hash_url(&url_clone));
            let cache_file = cache_dir.join(file_name);

            let mut bytes_opt = if let Ok(data) = std::fs::read(&cache_file) {
                Some(data)
            } else {
                None
            };

            if bytes_opt.is_none() {
                let mut req = client.get(&url_clone);
                if let Some(ref t) = token_opt {
                    req = req.header("X-Emby-Token", t);
                    req = req.header("Authorization", format!("MediaBrowser Token=\"{}\"", t));
                }
                if let Ok(resp) = req.send().await {
                    if resp.status().is_success() {
                        if let Ok(b) = resp.bytes().await {
                            let vec_bytes = b.to_vec();
                            let _ = std::fs::create_dir_all(cache_dir);
                            let _ = std::fs::write(&cache_file, &vec_bytes);
                            bytes_opt = Some(vec_bytes);
                        }
                    }
                }
            }

            let color_img = if let Some(bytes) = bytes_opt {
                tokio::task::spawn_blocking(move || {
                    if let Ok(dyn_img) = image::load_from_memory(&bytes) {
                        let rgba = dyn_img.to_rgba8();
                        let size = [rgba.width() as usize, rgba.height() as usize];
                        let pixels = rgba.into_raw();
                        Some(ColorImage::from_rgba_unmultiplied(size, &pixels))
                    } else {
                        None
                    }
                })
                .await
                .unwrap_or(None)
            } else {
                None
            };

            let _ = tx.send((url_clone, color_img));
            ctx_clone.request_repaint();
        });

        ImageStatus::Loading
    }
}