use crate::Result;
use ankit::AnkiClient;
use serde::Serialize;
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct EnrichQuery {
pub search: String,
pub empty_fields: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct EnrichCandidate {
pub note_id: i64,
pub model_name: String,
pub fields: HashMap<String, String>,
pub empty_fields: Vec<String>,
pub tags: Vec<String>,
}
#[derive(Debug, Clone, Default, Serialize)]
pub struct EnrichReport {
pub updated: usize,
pub failed: usize,
pub failures: Vec<EnrichFailure>,
}
#[derive(Debug, Clone, Serialize)]
pub struct EnrichFailure {
pub note_id: i64,
pub error: String,
}
#[derive(Debug)]
pub struct EnrichEngine<'a> {
client: &'a AnkiClient,
}
impl<'a> EnrichEngine<'a> {
pub(crate) fn new(client: &'a AnkiClient) -> Self {
Self { client }
}
pub async fn find_candidates(&self, query: &EnrichQuery) -> Result<Vec<EnrichCandidate>> {
let note_ids = self.client.notes().find(&query.search).await?;
if note_ids.is_empty() {
return Ok(Vec::new());
}
let note_infos = self.client.notes().info(¬e_ids).await?;
let mut candidates = Vec::new();
for info in note_infos {
let empty: Vec<String> = query
.empty_fields
.iter()
.filter(|field_name| {
info.fields
.get(*field_name)
.map(|f| f.value.trim().is_empty())
.unwrap_or(true) })
.cloned()
.collect();
if !empty.is_empty() {
let fields: HashMap<String, String> =
info.fields.into_iter().map(|(k, v)| (k, v.value)).collect();
candidates.push(EnrichCandidate {
note_id: info.note_id,
model_name: info.model_name,
fields,
empty_fields: empty,
tags: info.tags,
});
}
}
Ok(candidates)
}
pub async fn update_note(&self, note_id: i64, fields: &HashMap<String, String>) -> Result<()> {
self.client.notes().update_fields(note_id, fields).await?;
Ok(())
}
pub async fn update_notes(
&self,
updates: &[(i64, HashMap<String, String>)],
) -> Result<EnrichReport> {
let mut report = EnrichReport::default();
for (note_id, fields) in updates {
match self.client.notes().update_fields(*note_id, fields).await {
Ok(_) => report.updated += 1,
Err(e) => {
report.failed += 1;
report.failures.push(EnrichFailure {
note_id: *note_id,
error: e.to_string(),
});
}
}
}
Ok(report)
}
pub async fn tag_enriched(&self, note_ids: &[i64], tag: &str) -> Result<()> {
if !note_ids.is_empty() {
self.client.notes().add_tags(note_ids, tag).await?;
}
Ok(())
}
pub async fn pipeline(&self, query: &EnrichQuery) -> Result<EnrichmentPipeline> {
let candidates = self.find_candidates(query).await?;
Ok(EnrichmentPipeline::new(candidates))
}
}
#[derive(Debug, Clone)]
pub struct EnrichmentPipeline {
candidates: Vec<EnrichCandidate>,
updates: HashMap<i64, HashMap<String, String>>,
}
impl EnrichmentPipeline {
pub fn new(candidates: Vec<EnrichCandidate>) -> Self {
Self {
candidates,
updates: HashMap::new(),
}
}
pub fn candidates(&self) -> &[EnrichCandidate] {
&self.candidates
}
pub fn len(&self) -> usize {
self.candidates.len()
}
pub fn is_empty(&self) -> bool {
self.candidates.is_empty()
}
pub fn by_missing_field(&self) -> HashMap<String, Vec<&EnrichCandidate>> {
let mut groups: HashMap<String, Vec<&EnrichCandidate>> = HashMap::new();
for candidate in &self.candidates {
for field in &candidate.empty_fields {
groups.entry(field.clone()).or_default().push(candidate);
}
}
groups
}
pub fn by_model(&self) -> HashMap<String, Vec<&EnrichCandidate>> {
let mut groups: HashMap<String, Vec<&EnrichCandidate>> = HashMap::new();
for candidate in &self.candidates {
groups
.entry(candidate.model_name.clone())
.or_default()
.push(candidate);
}
groups
}
pub fn update(&mut self, note_id: i64, fields: HashMap<String, String>) {
self.updates.entry(note_id).or_default().extend(fields);
}
pub fn pending_updates(&self) -> usize {
self.updates.len()
}
pub fn pending_candidates(&self) -> Vec<&EnrichCandidate> {
self.candidates
.iter()
.filter(|c| !self.updates.contains_key(&c.note_id))
.collect()
}
pub async fn commit(&self, engine: &crate::Engine) -> Result<EnrichPipelineReport> {
let skipped = self
.candidates
.iter()
.filter(|c| !self.updates.contains_key(&c.note_id))
.count();
let mut updated = 0;
let mut failed = Vec::new();
for (note_id, fields) in &self.updates {
match engine.enrich().update_note(*note_id, fields).await {
Ok(_) => updated += 1,
Err(e) => {
failed.push((*note_id, e.to_string()));
}
}
}
Ok(EnrichPipelineReport {
updated,
failed,
skipped,
})
}
pub async fn commit_and_tag(
&self,
engine: &crate::Engine,
tag: &str,
) -> Result<EnrichPipelineReport> {
let report = self.commit(engine).await?;
if report.updated > 0 {
let updated_ids: Vec<i64> = self
.updates
.keys()
.filter(|id| !report.failed.iter().any(|(fid, _)| fid == *id))
.copied()
.collect();
if !updated_ids.is_empty() {
engine.enrich().tag_enriched(&updated_ids, tag).await?;
}
}
Ok(report)
}
}
#[derive(Debug, Clone, Default, Serialize)]
pub struct EnrichPipelineReport {
pub updated: usize,
pub failed: Vec<(i64, String)>,
pub skipped: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_enrich_query_construction() {
let query = EnrichQuery {
search: "deck:Test".to_string(),
empty_fields: vec!["Example".to_string(), "Audio".to_string()],
};
assert_eq!(query.search, "deck:Test");
assert_eq!(query.empty_fields.len(), 2);
assert!(query.empty_fields.contains(&"Example".to_string()));
}
#[test]
fn test_enrich_candidate_construction() {
let mut fields = HashMap::new();
fields.insert("Front".to_string(), "Hello".to_string());
fields.insert("Back".to_string(), "World".to_string());
let candidate = EnrichCandidate {
note_id: 12345,
model_name: "Basic".to_string(),
fields,
empty_fields: vec!["Example".to_string()],
tags: vec!["tag1".to_string()],
};
assert_eq!(candidate.note_id, 12345);
assert_eq!(candidate.model_name, "Basic");
assert_eq!(candidate.fields.len(), 2);
assert_eq!(candidate.empty_fields.len(), 1);
assert_eq!(candidate.tags.len(), 1);
}
#[test]
fn test_enrich_candidate_serialization() {
let candidate = EnrichCandidate {
note_id: 100,
model_name: "Vocab".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Definition".to_string()],
tags: vec![],
};
let json = serde_json::to_string(&candidate).unwrap();
assert!(json.contains("\"note_id\":100"));
assert!(json.contains("\"model_name\":\"Vocab\""));
}
#[test]
fn test_enrich_report_default() {
let report = EnrichReport::default();
assert_eq!(report.updated, 0);
assert_eq!(report.failed, 0);
assert!(report.failures.is_empty());
}
#[test]
fn test_enrich_report_construction() {
let failure = EnrichFailure {
note_id: 999,
error: "Not found".to_string(),
};
let report = EnrichReport {
updated: 5,
failed: 1,
failures: vec![failure],
};
assert_eq!(report.updated, 5);
assert_eq!(report.failed, 1);
assert_eq!(report.failures.len(), 1);
assert_eq!(report.failures[0].note_id, 999);
}
#[test]
fn test_enrich_failure_construction() {
let failure = EnrichFailure {
note_id: 12345,
error: "Field not found".to_string(),
};
assert_eq!(failure.note_id, 12345);
assert_eq!(failure.error, "Field not found");
}
#[test]
fn test_enrich_failure_serialization() {
let failure = EnrichFailure {
note_id: 456,
error: "Connection error".to_string(),
};
let json = serde_json::to_string(&failure).unwrap();
assert!(json.contains("\"note_id\":456"));
assert!(json.contains("\"error\":\"Connection error\""));
}
#[test]
fn test_enrichment_pipeline_new_empty() {
let pipeline = EnrichmentPipeline::new(vec![]);
assert!(pipeline.is_empty());
assert_eq!(pipeline.len(), 0);
assert_eq!(pipeline.pending_updates(), 0);
}
#[test]
fn test_enrichment_pipeline_with_candidates() {
let candidate = EnrichCandidate {
note_id: 1,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Back".to_string()],
tags: vec![],
};
let pipeline = EnrichmentPipeline::new(vec![candidate]);
assert!(!pipeline.is_empty());
assert_eq!(pipeline.len(), 1);
assert_eq!(pipeline.candidates().len(), 1);
}
#[test]
fn test_enrichment_pipeline_update() {
let candidate = EnrichCandidate {
note_id: 100,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Back".to_string()],
tags: vec![],
};
let mut pipeline = EnrichmentPipeline::new(vec![candidate]);
assert_eq!(pipeline.pending_updates(), 0);
let mut fields = HashMap::new();
fields.insert("Back".to_string(), "Answer".to_string());
pipeline.update(100, fields);
assert_eq!(pipeline.pending_updates(), 1);
}
#[test]
fn test_enrichment_pipeline_update_merge() {
let mut pipeline = EnrichmentPipeline::new(vec![]);
let mut fields1 = HashMap::new();
fields1.insert("Field1".to_string(), "Value1".to_string());
pipeline.update(100, fields1);
let mut fields2 = HashMap::new();
fields2.insert("Field2".to_string(), "Value2".to_string());
pipeline.update(100, fields2);
assert_eq!(pipeline.pending_updates(), 1);
}
#[test]
fn test_enrichment_pipeline_pending_candidates() {
let candidates = vec![
EnrichCandidate {
note_id: 1,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Back".to_string()],
tags: vec![],
},
EnrichCandidate {
note_id: 2,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Back".to_string()],
tags: vec![],
},
];
let mut pipeline = EnrichmentPipeline::new(candidates);
assert_eq!(pipeline.pending_candidates().len(), 2);
let mut fields = HashMap::new();
fields.insert("Back".to_string(), "Answer".to_string());
pipeline.update(1, fields);
assert_eq!(pipeline.pending_candidates().len(), 1);
assert_eq!(pipeline.pending_candidates()[0].note_id, 2);
}
#[test]
fn test_enrichment_pipeline_by_missing_field() {
let candidates = vec![
EnrichCandidate {
note_id: 1,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Field1".to_string()],
tags: vec![],
},
EnrichCandidate {
note_id: 2,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec!["Field1".to_string(), "Field2".to_string()],
tags: vec![],
},
];
let pipeline = EnrichmentPipeline::new(candidates);
let by_field = pipeline.by_missing_field();
assert_eq!(by_field.get("Field1").unwrap().len(), 2);
assert_eq!(by_field.get("Field2").unwrap().len(), 1);
}
#[test]
fn test_enrichment_pipeline_by_model() {
let candidates = vec![
EnrichCandidate {
note_id: 1,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec![],
tags: vec![],
},
EnrichCandidate {
note_id: 2,
model_name: "Cloze".to_string(),
fields: HashMap::new(),
empty_fields: vec![],
tags: vec![],
},
EnrichCandidate {
note_id: 3,
model_name: "Basic".to_string(),
fields: HashMap::new(),
empty_fields: vec![],
tags: vec![],
},
];
let pipeline = EnrichmentPipeline::new(candidates);
let by_model = pipeline.by_model();
assert_eq!(by_model.get("Basic").unwrap().len(), 2);
assert_eq!(by_model.get("Cloze").unwrap().len(), 1);
}
#[test]
fn test_enrich_pipeline_report_default() {
let report = EnrichPipelineReport::default();
assert_eq!(report.updated, 0);
assert!(report.failed.is_empty());
assert_eq!(report.skipped, 0);
}
#[test]
fn test_enrich_pipeline_report_construction() {
let report = EnrichPipelineReport {
updated: 10,
failed: vec![(100, "Error".to_string())],
skipped: 5,
};
assert_eq!(report.updated, 10);
assert_eq!(report.failed.len(), 1);
assert_eq!(report.skipped, 5);
}
#[test]
fn test_enrich_pipeline_report_serialization() {
let report = EnrichPipelineReport {
updated: 3,
failed: vec![],
skipped: 2,
};
let json = serde_json::to_string(&report).unwrap();
assert!(json.contains("\"updated\":3"));
assert!(json.contains("\"skipped\":2"));
}
}