use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
struct Stat {
mtime: i64,
size: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Processed {
mtime: i64,
size: u64,
content: String,
}
#[derive(Debug, Default, Serialize, Deserialize)]
struct QueueState {
pending: HashMap<String, Stat>,
processed: HashMap<String, Processed>,
}
fn key_of(path: &Path) -> String {
path.to_string_lossy().into_owned()
}
fn registry_key(index_dir: &Path) -> PathBuf {
PathBuf::from(crate::display_path(&index_dir.to_string_lossy()))
}
pub struct OcrQueue {
index_dir: PathBuf,
state: Mutex<QueueState>,
}
static REGISTRY: OnceLock<Mutex<HashMap<PathBuf, Arc<OcrQueue>>>> = OnceLock::new();
impl OcrQueue {
pub fn for_index_dir(index_dir: &Path) -> Arc<OcrQueue> {
let registry = REGISTRY.get_or_init(|| Mutex::new(HashMap::new()));
let key = registry_key(index_dir);
let mut map = registry.lock().expect("ocr queue registry mutex poisoned");
if let Some(existing) = map.get(&key) {
return existing.clone();
}
let state = load_state(index_dir);
let queue = Arc::new(OcrQueue {
index_dir: index_dir.to_path_buf(),
state: Mutex::new(state),
});
map.insert(key, queue.clone());
queue
}
pub(crate) fn cached_content(&self, path: &Path, mtime: i64, size: u64) -> Option<String> {
let state = self.state.lock().expect("ocr queue mutex poisoned");
state
.processed
.get(&key_of(path))
.filter(|p| p.mtime == mtime && p.size == size)
.map(|p| p.content.clone())
}
pub(crate) fn enqueue(&self, path: PathBuf, mtime: i64, size: u64) {
let mut state = self.state.lock().expect("ocr queue mutex poisoned");
state.pending.insert(key_of(&path), Stat { mtime, size });
}
pub(crate) fn pop_pending(&self) -> Option<(PathBuf, i64, u64)> {
let mut state = self.state.lock().expect("ocr queue mutex poisoned");
let key = state.pending.keys().next().cloned()?;
let stat = state.pending.remove(&key)?;
Some((PathBuf::from(key), stat.mtime, stat.size))
}
pub fn pending_len(&self) -> usize {
self.state
.lock()
.expect("ocr queue mutex poisoned")
.pending
.len()
}
pub(crate) fn mark_processed(&self, path: PathBuf, mtime: i64, size: u64, content: String) {
let mut state = self.state.lock().expect("ocr queue mutex poisoned");
state.processed.insert(
key_of(&path),
Processed {
mtime,
size,
content,
},
);
}
pub fn save(&self) -> Result<()> {
let state = self.state.lock().expect("ocr queue mutex poisoned");
save_state(&self.index_dir, &state)
}
pub fn compact(&self, roots: &[PathBuf]) {
let mut state = self.state.lock().expect("ocr queue mutex poisoned");
let keep = |key: &String| -> bool {
let path = Path::new(key);
roots.iter().any(|r| path.starts_with(r)) && path.exists()
};
state.pending.retain(|key, _| keep(key));
state.processed.retain(|key, _| keep(key));
}
}
fn queue_path(index_dir: &Path) -> PathBuf {
let stem = index_dir
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("dowse-index");
index_dir.with_file_name(format!("{stem}-ocr-queue.json"))
}
fn load_state(index_dir: &Path) -> QueueState {
let path = queue_path(index_dir);
match std::fs::read(&path) {
Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_else(|err| {
eprintln!(
"OCR 队列状态解析失败,按空队列处理 {}: {err}",
path.display()
);
QueueState::default()
}),
Err(_) => QueueState::default(),
}
}
fn save_state(index_dir: &Path, state: &QueueState) -> Result<()> {
let path = queue_path(index_dir);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let bytes = serde_json::to_vec(state)?;
std::fs::write(&path, bytes).context("写 OCR 队列状态失败")?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn save_and_load_state_round_trips_pending_and_processed() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
let mut state = QueueState::default();
state
.pending
.insert("C:\\shots\\a.png".to_string(), Stat { mtime: 1, size: 2 });
state.processed.insert(
"C:\\shots\\b.png".to_string(),
Processed {
mtime: 3,
size: 4,
content: "识别到的文字".to_string(),
},
);
save_state(&index_dir, &state).unwrap();
let loaded = load_state(&index_dir);
assert_eq!(loaded.pending.len(), 1);
assert_eq!(loaded.pending["C:\\shots\\a.png"].mtime, 1);
assert_eq!(loaded.processed["C:\\shots\\b.png"].content, "识别到的文字");
}
#[test]
fn load_state_missing_file_returns_empty_queue() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("no-such-index");
let state = load_state(&index_dir);
assert!(state.pending.is_empty());
assert!(state.processed.is_empty());
}
#[test]
fn for_index_dir_normalizes_extended_prefix_to_same_singleton() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
let prefixed = PathBuf::from(format!(r"\\?\{}", index_dir.display()));
let a = OcrQueue::for_index_dir(&index_dir);
let b = OcrQueue::for_index_dir(&prefixed);
assert!(
Arc::ptr_eq(&a, &b),
"两种路径写法应该命中同一个进程内单例,不能各建一份互相覆盖"
);
}
#[test]
fn enqueue_pop_and_mark_processed_round_trip() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
let queue = OcrQueue::for_index_dir(&index_dir);
let path = PathBuf::from("shot.png");
queue.enqueue(path.clone(), 10, 20);
assert_eq!(queue.pending_len(), 1);
assert!(queue.cached_content(&path, 10, 20).is_none());
let (popped_path, mtime, size) = queue.pop_pending().expect("队列非空应该能取出一项");
assert_eq!(popped_path, path);
assert_eq!((mtime, size), (10, 20));
assert_eq!(queue.pending_len(), 0, "取走之后队列应该清空");
queue.mark_processed(popped_path.clone(), mtime, size, "识别结果".to_string());
assert_eq!(
queue.cached_content(&popped_path, 10, 20).as_deref(),
Some("识别结果")
);
assert!(
queue.cached_content(&popped_path, 999, 20).is_none(),
"stat 对不上不该命中缓存"
);
}
#[test]
fn pending_progress_survives_save_and_reload_from_disk() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
{
let queue = OcrQueue::for_index_dir(&index_dir);
queue.enqueue(PathBuf::from("a.png"), 1, 2);
queue.enqueue(PathBuf::from("b.png"), 3, 4);
let (path, mtime, size) = queue.pop_pending().unwrap();
queue.mark_processed(path, mtime, size, "已识别".to_string());
queue.save().unwrap();
}
let reloaded = load_state(&index_dir);
assert_eq!(
reloaded.pending.len(),
1,
"还剩一张没处理,应该还在 pending 里"
);
assert_eq!(
reloaded.processed.len(),
1,
"已处理的那张应该被记住,重启不会重新识别"
);
}
#[test]
fn compact_drops_pending_outside_roots_and_missing_files() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
let root = dir.path().join("root");
std::fs::create_dir_all(&root).unwrap();
let kept = root.join("kept.png");
std::fs::write(&kept, b"x").unwrap();
let gone = root.join("gone.png");
std::fs::write(&gone, b"x").unwrap();
let outside = dir.path().join("outside.png");
std::fs::write(&outside, b"x").unwrap();
let queue = OcrQueue::for_index_dir(&index_dir);
queue.enqueue(kept.clone(), 1, 2);
queue.enqueue(gone.clone(), 1, 2);
queue.enqueue(outside.clone(), 1, 2);
std::fs::remove_file(&gone).unwrap();
queue.compact(std::slice::from_ref(&root));
assert_eq!(
queue.pending_len(),
1,
"只应保留 root 下、文件仍存在的 kept 一条"
);
}
#[test]
fn compact_drops_processed_cache_outside_roots_and_missing_files() {
let dir = tempfile::tempdir().unwrap();
let index_dir = dir.path().join("index");
let root = dir.path().join("root");
std::fs::create_dir_all(&root).unwrap();
let kept = root.join("kept.png");
std::fs::write(&kept, b"x").unwrap();
let gone = root.join("gone.png");
std::fs::write(&gone, b"x").unwrap();
let outside = dir.path().join("outside.png");
std::fs::write(&outside, b"x").unwrap();
let queue = OcrQueue::for_index_dir(&index_dir);
queue.mark_processed(kept.clone(), 1, 2, "内容".to_string());
queue.mark_processed(gone.clone(), 1, 2, "内容".to_string());
queue.mark_processed(outside.clone(), 1, 2, "内容".to_string());
std::fs::remove_file(&gone).unwrap();
queue.compact(std::slice::from_ref(&root));
assert!(
queue.cached_content(&kept, 1, 2).is_some(),
"root 下且文件仍在应保留缓存"
);
assert!(
queue.cached_content(&gone, 1, 2).is_none(),
"文件已删应丢弃缓存"
);
assert!(
queue.cached_content(&outside, 1, 2).is_none(),
"root 之外应丢弃缓存"
);
}
}