src/batch.rs
Rama main · réplica de lectura del repositorio original.
//! 批量下载编排：勾选歌曲 -> 逐首 搜索 + 匹配 + 下载
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
use regex::Regex;
use serde::Serialize;
use serde_json::{json, Value};
use crate::config::Settings;
use crate::music::{self, SongView};
use crate::state::{gen_id, now_ts, TaskStore};
fn norm(s: &str) -> String {
static RE: OnceLock<Regex> = OnceLock::new();
let re = RE.get_or_init(|| Regex::new(r#"[\s()（）\[\]【】\-_.·～~']"#).unwrap());
re.replace_all(&s.to_lowercase(), "").into_owned()
}
/// Ratcliff-Obershelp 最长公共子串递归匹配数（对齐 difflib.SequenceMatcher）
fn match_count(a: &[char], b: &[char]) -> usize {
if a.is_empty() || b.is_empty() {
return 0;
}
let mut dp = vec![vec![0usize; b.len() + 1]; a.len() + 1];
let (mut best, mut bi, mut bj) = (0, 0, 0);
for i in 1..=a.len() {
for j in 1..=b.len() {
if a[i - 1] == b[j - 1] {
dp[i][j] = dp[i - 1][j - 1] + 1;
if dp[i][j] > best {
best = dp[i][j];
bi = i - best;
bj = j - best;
}
}
}
}
if best == 0 {
0
} else {
best + match_count(&a[..bi], &b[..bj]) + match_count(&a[bi + best..], &b[bj + best..])
}
}
fn sequence_ratio(a: &str, b: &str) -> f64 {
let a: Vec<char> = a.chars().collect();
let b: Vec<char> = b.chars().collect();
let total = a.len() + b.len();
if total == 0 {
return 1.0;
}
2.0 * match_count(&a, &b) as f64 / total as f64
}
/// 候选歌曲与期望的相似度打分：标题 0.6 + 歌手 0.3 + 无损 0.1
pub fn match_score(title: &str, artist: &str, cand: &SongView) -> f64 {
let (t1, t2) = (norm(title), norm(&cand.title));
let (a1, a2) = (norm(artist), norm(&cand.artist));
if t1.is_empty() || t2.is_empty() {
return 0.0;
}
let ts = if t1.contains(&t2) || t2.contains(&t1) {
1.0
} else {
sequence_ratio(&t1, &t2)
};
let as_ = if !a1.is_empty() && !a2.is_empty() {
if a1.contains(&a2) || a2.contains(&a1) {
1.0
} else {
sequence_ratio(&a1, &a2)
}
} else {
0.0
};
let quality = if cand.ext == "flac" { 0.1 } else { 0.0 };
ts * 0.6 + as_ * 0.3 + quality
}
/// 批量下载批次状态（内存实现）
#[derive(Default)]
pub struct BatchStore {
batches: Mutex<HashMap<String, Batch>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct Batch {
pub id: String,
pub created_at: f64,
pub songs: Vec<BatchSong>,
}
#[derive(Debug, Clone, Serialize)]
pub struct BatchSong {
pub title: String,
pub artist: String,
pub status: String, // pending / matching / downloading / done / failed / unmatched
pub source: Option<String>,
pub ext: Option<String>,
pub file: Option<String>,
pub error: Option<String>,
}
impl BatchStore {
pub fn new() -> Self {
Self::default()
}
pub fn create(&self, songs: Vec<(String, String)>) -> String {
let batch_id = gen_id(12);
let batch = Batch {
id: batch_id.clone(),
created_at: now_ts(),
songs: songs
.into_iter()
.map(|(title, artist)| BatchSong {
title,
artist,
status: "pending".to_string(),
source: None,
ext: None,
file: None,
error: None,
})
.collect(),
};
self.batches
.lock()
.unwrap()
.insert(batch_id.clone(), batch);
batch_id
}
pub fn update_song(&self, batch_id: &str, idx: usize, f: impl FnOnce(&mut BatchSong)) {
if let Some(batch) = self.batches.lock().unwrap().get_mut(batch_id) {
if let Some(song) = batch.songs.get_mut(idx) {
f(song);
}
}
}
pub fn get(&self, batch_id: &str) -> Option<Value> {
let batch = self.batches.lock().unwrap().get(batch_id).cloned()?;
let mut summary: HashMap<&str, usize> = HashMap::new();
for s in &batch.songs {
*summary.entry(s.status.as_str()).or_insert(0) += 1;
}
Some(json!({
"id": batch.id,
"created_at": batch.created_at,
"songs": batch.songs,
"summary": summary,
}))
}
}
/// 串行 worker：单首歌 搜索 -> 打分匹配 -> 下载（状态同步到批次与任务中心）
#[allow(clippy::too_many_arguments)]
pub async fn run_batch_song(
cfg: Arc<Settings>,
batches: Arc<BatchStore>,
tasks: Arc<TaskStore>,
batch_id: String,
idx: usize,
title: String,
artist: String,
task_id: String,
) {
batches.update_song(&batch_id, idx, |s| s.status = "matching".into());
tasks.update(&task_id, |t| t.status = "matching".into());
let keyword = format!("{artist} {title}").trim().to_string();
let outcome = music::search(cfg.clone(), keyword).await;
let mut best_idx: Option<usize> = None;
let mut best_score = 0.0f64;
for (i, (_, view, _)) in outcome.items.iter().enumerate() {
let score = match_score(&title, &artist, view);
if score > best_score {
best_score = score;
best_idx = Some(i);
}
}
let Some(best_i) = best_idx.filter(|_| best_score >= cfg.match_threshold) else {
let msg = format!("最高匹配度 {best_score:.2}");
batches.update_song(&batch_id, idx, |s| {
s.status = "unmatched".into();
s.error = Some(msg.clone());
});
tasks.update(&task_id, |t| {
t.status = "unmatched".into();
t.error = Some(msg);
});
return;
};
// 提前取出 owned 值，避免引用跨越 await
let (_, view, raw) = &outcome.items[best_i];
let (source_name, ext, raw) = (view.source.clone(), view.ext.clone(), raw.clone());
batches.update_song(&batch_id, idx, |s| {
s.status = "downloading".into();
s.source = Some(source_name.clone());
s.ext = Some(ext.clone());
});
tasks.update(&task_id, |t| {
t.status = "downloading".into();
t.source = Some(source_name);
t.ext = Some(ext);
});
match music::download(cfg, raw).await {
Ok(path) => {
let file = path.display().to_string();
batches.update_song(&batch_id, idx, |s| {
s.status = "done".into();
s.file = Some(file.clone());
});
tasks.update(&task_id, |t| {
t.status = "done".into();
t.file = Some(file.clone());
});
tracing::info!("batch matched: {title} -> {file} ({best_score:.2})");
}
Err(e) => {
batches.update_song(&batch_id, idx, |s| {
s.status = "failed".into();
s.error = Some(e.clone());
});
tasks.update(&task_id, |t| {
t.status = "failed".into();
t.error = Some(e);
});
}
}
}