use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{mpsc, Arc, Condvar, Mutex};
use std::time::Duration;
use zune_jpeg::zune_core::bytestream::ZCursor;
use zune_jpeg::zune_core::colorspace::ColorSpace;
use zune_jpeg::zune_core::options::DecoderOptions;
use zune_jpeg::JpegDecoder;
use crate::library::{self, Kind};
use crate::mpv::{Event, Mpv};
use crate::vk::{Thumb, THUMB_H, THUMB_SLOTS, THUMB_W};
const WORKERS: usize = 4;
const LOCAL_WORKERS: usize = 2;
const LOCAL_TIMEOUT: Duration = Duration::from_secs(15);
const COVER_STEMS: &[&str] = &["cover", "folder", "poster", "front", "album"];
const IMAGE_EXT: &[&str] = &["jpg", "jpeg", "png", "webp"];
const QUEUE_LIMIT: usize = 48;
const UPLOADS_PER_PUMP: usize = 8;
const TTL: Duration = Duration::from_secs(30 * 24 * 3600);
pub struct Image {
pub w: u32,
pub h: u32,
pub rgba: Vec<u8>,
}
enum State {
Queued,
Ready(Thumb),
Failed,
}
type Queue = Arc<(Mutex<Vec<String>>, Condvar)>;
pub struct Thumbs {
states: HashMap<String, State>,
queue: Queue, local_queue: Queue, done: mpsc::Receiver<(String, Option<Image>)>,
slots: Vec<Option<(String, u64)>>, frame: u64,
}
impl Thumbs {
pub fn new() -> Thumbs {
let queue: Queue = Arc::new((Mutex::new(Vec::new()), Condvar::new()));
let local_queue: Queue = Arc::new((Mutex::new(Vec::new()), Condvar::new()));
let (tx, done) = mpsc::channel();
let agent: ureq::Agent = ureq::Agent::config_builder()
.timeout_global(Some(Duration::from_secs(10)))
.build()
.into();
for _ in 0..WORKERS {
let (queue, tx, agent) = (queue.clone(), tx.clone(), agent.clone());
std::thread::spawn(move || worker(&queue, &tx, |url| load(&agent, url)));
}
for _ in 0..LOCAL_WORKERS {
let (queue, tx) = (local_queue.clone(), tx.clone());
std::thread::spawn(move || worker(&queue, &tx, |path| load_local(Path::new(path))));
}
Thumbs { states: HashMap::new(), queue, local_queue, done, slots: vec![None; THUMB_SLOTS], frame: 0 }
}
pub fn begin_frame(&mut self) {
self.frame += 1;
}
pub fn get(&mut self, url: &str) -> Option<Thumb> {
match self.states.get(url) {
Some(State::Ready(thumb)) => {
let thumb = *thumb;
if let Some(slot) = self.slots.iter_mut().flatten().find(|(owner, _)| owner == url) {
slot.1 = self.frame;
}
Some(thumb)
}
Some(_) => None,
None => {
self.request(url);
None
}
}
}
pub fn prefetch(&mut self, url: &str) {
if !self.states.contains_key(url) {
self.request(url);
}
}
fn request(&mut self, url: &str) {
if url.is_empty() {
return;
}
self.states.insert(url.to_string(), State::Queued);
let queue = if url.starts_with('/') { &self.local_queue } else { &self.queue };
let (lock, ready) = &**queue;
let mut queue = lock.lock().unwrap();
queue.push(url.to_string());
if queue.len() > QUEUE_LIMIT {
let dropped = queue.remove(0);
self.states.remove(&dropped);
}
ready.notify_one();
}
pub fn pump(&mut self, mut upload: impl FnMut(usize, u32, u32, &[u8]) -> Result<Thumb, String>) -> bool {
let mut changed = false;
for (url, image) in self.done.try_iter().take(UPLOADS_PER_PUMP).collect::<Vec<_>>() {
let Some(image) = image else {
self.states.insert(url, State::Failed);
continue;
};
let Some(slot) = self.free_slot() else {
self.states.remove(&url); continue;
};
match upload(slot, image.w, image.h, &image.rgba) {
Ok(thumb) => {
self.slots[slot] = Some((url.clone(), self.frame));
self.states.insert(url, State::Ready(thumb));
changed = true;
}
Err(error) => {
eprintln!("thumbs: {error}");
self.states.insert(url, State::Failed);
}
}
}
changed
}
fn free_slot(&mut self) -> Option<usize> {
if let Some(i) = self.slots.iter().position(Option::is_none) {
return Some(i);
}
let (i, _) = self
.slots
.iter()
.enumerate()
.filter_map(|(i, s)| s.as_ref().map(|(_, used)| (i, *used)))
.filter(|&(_, used)| used < self.frame)
.min_by_key(|&(_, used)| used)?;
if let Some((owner, _)) = self.slots[i].take() {
self.states.remove(&owner);
}
Some(i)
}
}
fn worker(queue: &Queue, tx: &mpsc::Sender<(String, Option<Image>)>, load: impl Fn(&str) -> Option<Image>) {
let (lock, ready) = &**queue;
loop {
let url = {
let mut pending = lock.lock().unwrap();
loop {
if let Some(url) = pending.pop() {
break url;
}
pending = ready.wait(pending).unwrap();
}
};
let image = load(&url);
if tx.send((url, image)).is_err() {
return;
}
}
}
fn load(agent: &ureq::Agent, url: &str) -> Option<Image> {
let path = cache_path(url);
let cached = path.as_ref().and_then(|p| {
let age = std::fs::metadata(p).ok()?.modified().ok()?.elapsed().ok()?;
(age < TTL).then(|| std::fs::read(p).ok()).flatten()
});
let bytes = match cached {
Some(bytes) => bytes,
None => {
let mut resp = agent.get(url).call().map_err(|e| eprintln!("thumbs: {url}: {e}")).ok()?;
let bytes = resp.body_mut().read_to_vec().ok()?;
if let Some(p) = &path {
let _ = p.parent().map(std::fs::create_dir_all);
let _ = std::fs::write(p, &bytes);
}
bytes
}
};
decode(&bytes, url.contains("i.ytimg.com/vi/"))
}
fn cache_dir() -> Option<PathBuf> {
Some(PathBuf::from(std::env::var("HOME").ok()?).join("medias/.cache/thumbs"))
}
fn fnv1a(bytes: &[u8]) -> u64 {
bytes.iter().fold(0xcbf29ce484222325u64, |hash, &b| (hash ^ b as u64).wrapping_mul(0x100000001b3))
}
fn cache_path(url: &str) -> Option<PathBuf> {
Some(cache_dir()?.join(format!("{:016x}.jpg", fnv1a(url.as_bytes()))))
}
fn local_cache_path(path: &Path, meta: &std::fs::Metadata) -> Option<PathBuf> {
let mtime = meta.modified().ok()?.duration_since(std::time::UNIX_EPOCH).ok()?.as_secs();
let key = format!("{}\0{mtime}\0{}", path.display(), meta.len());
Some(cache_dir()?.join(format!("{:016x}.jpg", fnv1a(key.as_bytes()))))
}
fn load_local(path: &Path) -> Option<Image> {
let meta = std::fs::metadata(path).ok()?;
let cache = local_cache_path(path, &meta)?;
let none = cache.with_extension("none");
if let Ok(bytes) = std::fs::read(&cache) {
return decode(&bytes, false);
}
if none.exists() {
return None;
}
let source = if meta.is_dir() { dir_source(path, 2) } else { Some(path.to_path_buf()) };
let image = source.and_then(|source| {
let bytes = if is_jpeg(&source) { std::fs::read(&source).ok()? } else { render_frame(&source)? };
let image = decode(&bytes, false)?;
let _ = std::fs::create_dir_all(cache.parent()?);
let _ = std::fs::write(&cache, &bytes);
Some(image)
});
if image.is_none() {
let _ = std::fs::create_dir_all(none.parent()?);
let _ = std::fs::write(&none, b"");
}
image
}
fn is_jpeg(path: &Path) -> bool {
path.extension().and_then(|e| e.to_str()).is_some_and(|e| e.eq_ignore_ascii_case("jpg") || e.eq_ignore_ascii_case("jpeg"))
}
fn dir_source(dir: &Path, depth: usize) -> Option<PathBuf> {
let mut covers: Vec<(usize, PathBuf)> = std::fs::read_dir(dir)
.ok()?
.flatten()
.map(|e| e.path())
.filter_map(|p| {
let ext = p.extension()?.to_str()?.to_ascii_lowercase();
let stem = p.file_stem()?.to_str()?.to_ascii_lowercase();
let rank = COVER_STEMS.iter().position(|s| *s == stem)?;
IMAGE_EXT.contains(&ext.as_str()).then_some((rank, p))
})
.collect();
covers.sort();
if let Some((_, cover)) = covers.into_iter().next() {
return Some(cover);
}
let entries = library::list(dir).ok()?;
if let Some(media) = entries.iter().find(|e| e.kind != Kind::Dir) {
return Some(media.path.clone());
}
let sub = entries.iter().find(|e| e.kind == Kind::Dir)?;
(depth > 0).then(|| dir_source(&sub.path, depth - 1)).flatten()
}
fn render_frame(source: &Path) -> Option<Vec<u8>> {
static RUN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let n = RUN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let out = cache_dir()?.join(format!("tmp-{}-{n}", std::process::id()));
std::fs::create_dir_all(&out).ok()?;
let frame = (|| {
let m = Mpv::new().ok()?;
let scale = format!("scale=w={THUMB_W}:h={THUMB_H}:force_original_aspect_ratio=decrease");
let out_dir = out.to_string_lossy();
for (name, value) in [
("vo", "image"), ("vo-image-format", "jpg"), ("vo-image-jpeg-quality", "85"),
("vo-image-outdir", &out_dir), ("vf", &scale), ("frames", "1"),
("start", "15%"), ("hr-seek", "no"), ("image-display-duration", "0"),
("ao", "null"), ("aid", "no"), ("sid", "no"), ("hwdec", "no"),
("idle", "no"), ("load-scripts", "no"), ("ytdl", "no"), ("msg-level", "all=no"),
] {
m.set_option(name, value).ok()?;
}
m.init().ok()?;
m.loadfile(&source.to_string_lossy()).ok()?;
let start = std::time::Instant::now();
loop {
match m.wait_event(0.25) {
Event::EndFile(_) | Event::PlaybackError | Event::Shutdown => break,
_ if start.elapsed() > LOCAL_TIMEOUT => break,
_ => {}
}
}
drop(m);
std::fs::read(out.join("00000001.jpg")).ok()
})();
let _ = std::fs::remove_dir_all(&out);
frame
}
pub fn decode(bytes: &[u8], letterboxed: bool) -> Option<Image> {
let options = DecoderOptions::default().jpeg_set_out_colorspace(ColorSpace::RGBA);
let mut decoder = JpegDecoder::new_with_options(ZCursor::new(bytes), options);
let pixels = decoder.decode().ok()?;
let info = decoder.info()?;
let (w, h) = (info.width as usize, info.height as usize);
if w == 0 || h == 0 {
return None;
}
let rgba = match pixels.len() / (w * h) {
4 => pixels,
3 => pixels.chunks_exact(3).flat_map(|p| [p[0], p[1], p[2], 255]).collect(),
1 => pixels.iter().flat_map(|&l| [l, l, l, 255]).collect(),
_ => return None,
};
let (y0, ch) = if letterboxed && w * 3 == h * 4 { ((h - w * 9 / 16) / 2, w * 9 / 16) } else { (0, h) };
Some(fit(&rgba, w, y0, ch, THUMB_W as usize, THUMB_H as usize))
}
fn fit(rgba: &[u8], w: usize, y0: usize, ch: usize, max_w: usize, max_h: usize) -> Image {
let scale = (max_w as f64 / w as f64).min(max_h as f64 / ch as f64).min(1.0);
let tw = ((w as f64 * scale).round() as usize).clamp(1, max_w);
let th = ((ch as f64 * scale).round() as usize).clamp(1, max_h);
let mut out = Vec::with_capacity(tw * th * 4);
for ty in 0..th {
let (sy0, sy1) = (ty * ch / th, ((ty + 1) * ch / th).max(ty * ch / th + 1));
for tx in 0..tw {
let (sx0, sx1) = (tx * w / tw, ((tx + 1) * w / tw).max(tx * w / tw + 1));
let mut sum = [0u32; 4];
for sy in sy0..sy1 {
let row = (y0 + sy) * w;
for sx in sx0..sx1 {
let p = &rgba[(row + sx) * 4..(row + sx) * 4 + 4];
for c in 0..4 {
sum[c] += p[c] as u32;
}
}
}
let n = ((sy1 - sy0) * (sx1 - sx0)) as u32;
out.extend(sum.map(|s| ((s + n / 2) / n) as u8));
}
}
Image { w: tw as u32, h: th as u32, rgba: out }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fit_box_filters_and_keeps_aspect() {
let mut rgba = Vec::new();
for _ in 0..2 {
for x in 0..4 {
let v = if x < 2 { 0 } else { 255 };
rgba.extend([v, v, v, 255]);
}
}
let image = fit(&rgba, 4, 0, 2, 2, 2);
assert_eq!((image.w, image.h), (2, 1));
assert_eq!(image.rgba, [0, 0, 0, 255, 255, 255, 255, 255]);
let image = fit(&rgba, 4, 0, 2, 480, 270);
assert_eq!((image.w, image.h, image.rgba.len()), (4, 2, 32));
let image = fit(&rgba, 4, 1, 1, 480, 270);
assert_eq!((image.w, image.h), (4, 1));
}
#[test]
fn cache_paths_are_stable_and_distinct() {
let a = cache_path("https://i.ytimg.com/vi/a/hqdefault.jpg");
let b = cache_path("https://i.ytimg.com/vi/b/hqdefault.jpg");
assert_eq!(a, cache_path("https://i.ytimg.com/vi/a/hqdefault.jpg"));
assert_ne!(a, b);
}
#[test]
fn directories_prefer_cover_files_then_first_media() {
let root = std::env::temp_dir().join(format!("medias-thumbs-{}", std::process::id()));
let album = root.join("Artist/Album");
std::fs::create_dir_all(&album).unwrap();
std::fs::write(album.join("02 b.flac"), b"").unwrap();
std::fs::write(album.join("01 a.flac"), b"").unwrap();
assert_eq!(dir_source(&album, 0), Some(album.join("01 a.flac")));
assert_eq!(dir_source(&root, 2), Some(album.join("01 a.flac"))); assert_eq!(dir_source(&root, 1), None);
std::fs::write(album.join("Folder.PNG"), b"").unwrap();
assert_eq!(dir_source(&album, 0), Some(album.join("Folder.PNG")));
std::fs::write(album.join("cover.jpg"), b"").unwrap();
assert_eq!(dir_source(&album, 0), Some(album.join("cover.jpg"))); std::fs::remove_dir_all(&root).unwrap();
}
#[test]
#[ignore]
fn hqdefault_decodes_to_a_full_slot() {
let agent: ureq::Agent = ureq::Agent::config_builder().build().into();
let mut resp = agent.get("https://i.ytimg.com/vi/V_uq3v6xK-4/hqdefault.jpg").call().unwrap();
let bytes = resp.body_mut().read_to_vec().unwrap();
let image = decode(&bytes, true).unwrap();
assert_eq!((image.w, image.h), (THUMB_W, THUMB_H));
assert_eq!(image.rgba.len(), (THUMB_W * THUMB_H * 4) as usize);
}
#[test]
#[ignore]
fn local_frames_render_through_libmpv() {
let dir = std::env::temp_dir().join(format!("medias-frames-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let ffmpeg = |args: &[&str]| {
assert!(std::process::Command::new("ffmpeg").args(["-v", "error", "-y"]).args(args).status().unwrap().success())
};
let (video, art, plain) = (dir.join("v.mkv"), dir.join("a.mp3"), dir.join("p.mp3"));
ffmpeg(&["-f", "lavfi", "-i", "testsrc2=size=1920x1080:rate=25:duration=10", "-c:v", "libx264", "-preset", "ultrafast", video.to_str().unwrap()]);
ffmpeg(&["-f", "lavfi", "-i", "sine=duration=5", plain.to_str().unwrap()]);
ffmpeg(&["-f", "lavfi", "-i", "testsrc2=size=600x600:duration=1", "-frames:v", "1", dir.join("c.jpg").to_str().unwrap()]);
ffmpeg(&["-i", plain.to_str().unwrap(), "-i", dir.join("c.jpg").to_str().unwrap(), "-map", "0", "-map", "1",
"-c", "copy", "-disposition:v", "attached_pic", art.to_str().unwrap()]);
std::fs::remove_file(dir.join("c.jpg")).unwrap(); let frame = decode(&render_frame(&video).unwrap(), false).unwrap();
assert_eq!((frame.w, frame.h), (THUMB_W, THUMB_H));
let cover = decode(&render_frame(&art).unwrap(), false).unwrap();
assert_eq!((cover.w, cover.h), (THUMB_H, THUMB_H));
assert!(render_frame(&plain).is_none());
std::fs::remove_dir_all(&dir).unwrap();
}
}