use std::collections::{HashMap, HashSet};
use std::path::Path;
use uuid::Uuid;
use crate::ai::AiClient;
use crate::config::Config;
use crate::drift::{
assemble_descriptions, attribute_continuations, parse_drift_pairs, resolve_conflicts, Candidate,
DescriptionSnippet, DriftReport, EntityDescriptions, EntityKind,
};
use crate::error::{Error, Result};
use crate::project::ProjectLayout;
use crate::store::hierarchy::Hierarchy;
use crate::store::node::NodeKind;
use crate::store::{Store, SYSTEM_TAG_ARTEFACTS, SYSTEM_TAG_CHARACTERS, SYSTEM_TAG_PLACES};
use super::DriftCommand;
const DRIFT_SYSTEM_PROMPT: &str = "You are a continuity editor for a work of fiction. You receive \
NUMBERED descriptions of a SINGLE entity (a character, place, or object), each drawn from a \
different point in the manuscript, in chapter order. Flag pairs that CONTRADICT each other — the \
same attribute described in incompatible ways (a place cramped vs spacious, smoky vs airy; a \
character soft-spoken vs booming; an object pristine vs battered) with no in-story event that \
would explain the change. Do NOT flag descriptions that merely add new detail, describe different \
aspects, or reflect a change the story clearly dramatizes. Output ONE contradiction per line, in \
the exact form:\n\
i | j | why\n\
where `i` and `j` are the description NUMBERS and `why` is a one-line explanation of the \
contradiction. Output nothing else — no preamble, no commentary, no markdown. If the descriptions \
are consistent, output nothing.";
pub fn run(project: &Path, cmd: DriftCommand) -> Result<()> {
match cmd {
DriftCommand::List { json, entity } => list(project, json, entity.as_deref()),
DriftCommand::Scan { provider, json } => scan(project, provider.as_deref(), json),
}
}
pub fn collect_entity_descriptions(project: &Path) -> Result<Vec<EntityDescriptions>> {
let layout = ProjectLayout::new(project);
let cfg = Config::load_layered(&layout.config_path())?;
let store = Store::open(layout.clone(), &cfg).map_err(|e| Error::Store(e.to_string()))?;
let hierarchy = Hierarchy::load(&store).map_err(|e| Error::Store(e.to_string()))?;
Ok(gather(&store, &hierarchy, &cfg.drift, &cfg.language))
}
fn gather(
store: &Store,
hierarchy: &Hierarchy,
cfg: &crate::config::DriftConfig,
language: &str,
) -> Vec<EntityDescriptions> {
let index = chapter_index(hierarchy);
let lexicon = entities(hierarchy);
let chapters = chapter_paragraphs(store, hierarchy);
let coref = attribute_continuations(&chapters, &lexicon, language);
let mut out = Vec::new();
for (entity, kind) in lexicon.iter().cloned() {
let coref_ids: HashSet<Uuid> = coref
.iter()
.filter(|(_, names)| names.iter().any(|n| n.eq_ignore_ascii_case(&entity)))
.map(|(p, _)| *p)
.collect();
let snippets = retrieve(store, &index, &entity, cfg, &coref_ids);
if !snippets.is_empty() {
out.push(EntityDescriptions { entity, kind, snippets });
}
}
out
}
fn chapter_paragraphs(store: &Store, h: &Hierarchy) -> Vec<Vec<(Uuid, String)>> {
let mut chapters = Vec::new();
for book in h.iter().filter(|n| n.kind == NodeKind::Book && n.system_tag.is_none()) {
for chapter in h.children_of(Some(book.id)) {
if chapter.kind != NodeKind::Chapter {
continue;
}
let mut paras = Vec::new();
for pid in h.collect_subtree(chapter.id) {
if h.get(pid).map(|n| n.kind) != Some(NodeKind::Paragraph) {
continue;
}
if let Ok(Some(bytes)) = store.get_content(pid) {
let text = crate::audiobook::typst_to_plain(&String::from_utf8_lossy(&bytes))
.trim()
.to_string();
if !text.is_empty() {
paras.push((pid, text));
}
}
}
if !paras.is_empty() {
chapters.push(paras);
}
}
}
chapters
}
fn chapter_index(h: &Hierarchy) -> HashMap<Uuid, (usize, String)> {
let mut map = HashMap::new();
let mut order = 0usize;
for book in h.iter().filter(|n| n.kind == NodeKind::Book && n.system_tag.is_none()) {
for chapter in h.children_of(Some(book.id)) {
if chapter.kind != NodeKind::Chapter {
continue;
}
let title = if chapter.title.trim().is_empty() {
chapter.slug.clone()
} else {
chapter.title.clone()
};
for pid in h.collect_subtree(chapter.id) {
if h.get(pid).map(|n| n.kind) == Some(NodeKind::Paragraph) {
map.insert(pid, (order, title.clone()));
}
}
order += 1;
}
}
map
}
fn entities(h: &Hierarchy) -> Vec<(String, EntityKind)> {
entities_with_nodes(h).into_iter().map(|(n, k, _)| (n, k)).collect()
}
pub fn entities_with_nodes(h: &Hierarchy) -> Vec<(String, EntityKind, Uuid)> {
let books = [
(SYSTEM_TAG_CHARACTERS, EntityKind::Character),
(SYSTEM_TAG_PLACES, EntityKind::Place),
(SYSTEM_TAG_ARTEFACTS, EntityKind::Artefact),
];
let mut out = Vec::new();
for (tag, kind) in books {
let Some(book) = h
.iter()
.find(|n| n.kind == NodeKind::Book && n.system_tag.as_deref() == Some(tag))
else {
continue;
};
for id in h.collect_subtree(book.id) {
if let Some(n) = h.get(id) {
if n.kind == NodeKind::Paragraph && !n.title.trim().is_empty() {
out.push((n.title.trim().to_string(), kind, n.id));
}
}
}
}
out
}
fn retrieve(
store: &Store,
index: &HashMap<Uuid, (usize, String)>,
entity: &str,
cfg: &crate::config::DriftConfig,
coref_ids: &HashSet<Uuid>,
) -> Vec<DescriptionSnippet> {
let query = format!("{entity} description appearance manner voice condition");
let raw = match store.search_text(&query, cfg.top_k) {
Ok(r) => r,
Err(_) => return Vec::new(),
};
let mut candidates = Vec::new();
for v in raw {
let Some(id) = v
.get("id")
.and_then(|x| x.as_str())
.and_then(|s| Uuid::parse_str(s).ok())
else {
continue;
};
let Some((order, title)) = index.get(&id) else {
continue; };
if let Ok(Some(bytes)) = store.get_content(id) {
let text = crate::audiobook::typst_to_plain(&String::from_utf8_lossy(&bytes))
.trim()
.to_string();
if text.is_empty() {
continue;
}
candidates.push(Candidate {
paragraph: id,
chapter_order: *order,
chapter_title: title.clone(),
text,
});
}
}
assemble_descriptions(entity, &candidates, cfg.max_snippets, coref_ids)
}
fn list(project: &Path, json: bool, entity: Option<&str>) -> Result<()> {
let mut descs = collect_entity_descriptions(project)?;
if let Some(name) = entity {
let needle = name.to_lowercase();
descs.retain(|d| d.entity.to_lowercase().contains(&needle));
}
if json {
let payload = serde_json::to_string_pretty(&descs)
.map_err(|e| Error::Store(format!("serialize drift descriptions: {e}")))?;
println!("{payload}");
return Ok(());
}
if descs.is_empty() {
println!(
"drift: no entity descriptions retrieved — populate the Characters / Places / \
Artefacts books, and make sure the vector index is built (open + save once, or \
reindex)."
);
return Ok(());
}
let total: usize = descs.iter().map(|d| d.snippets.len()).sum();
println!(
"drift: {} entit{} described across {total} paragraph(s)\n",
descs.len(),
if descs.len() == 1 { "y" } else { "ies" }
);
for d in &descs {
println!("{} ({}) — {} snippet(s):", d.entity, d.kind.label(), d.snippets.len());
for s in &d.snippets {
let preview: String = s.text.chars().take(100).collect();
let ell = if s.text.chars().count() > 100 { "…" } else { "" };
println!(" · [{}] {preview}{ell}", s.chapter);
}
println!();
}
Ok(())
}
fn scan(project: &Path, provider: Option<&str>, json: bool) -> Result<()> {
let layout = ProjectLayout::new(project);
layout.require_initialized()?;
let cfg = Config::load_layered(&layout.config_path())?;
let store = Store::open(layout.clone(), &cfg).map_err(|e| Error::Store(e.to_string()))?;
let hierarchy = Hierarchy::load(&store).map_err(|e| Error::Store(e.to_string()))?;
let report = scan_with(&store, &hierarchy, &cfg, &layout, provider, &|s| eprintln!("{s}"))?;
if json {
let rendered = serde_json::to_string_pretty(&report)
.map_err(|e| Error::Store(format!("drift JSON: {e}")))?;
println!("{rendered}");
} else if report.conflicts.is_empty() {
let n = report.descriptions.iter().filter(|d| d.snippets.len() >= 2).count();
println!(
"drift scan: ✓ no description contradictions across {n} entit{}",
if n == 1 { "y" } else { "ies" }
);
} else {
println!("drift scan: {} description contradiction(s):", report.conflicts.len());
for c in &report.conflicts {
println!(
" ⚠ {} ({}) — [{}] “{}” ⟷ [{}] “{}”\n ↳ {}",
c.entity, c.kind.label(), c.chapter_a, c.a, c.chapter_b, c.b, c.detail
);
}
eprintln!(" (also surfaced in `inkhaven edit`)");
}
Ok(())
}
pub fn scan_with(
store: &Store,
hierarchy: &Hierarchy,
cfg: &Config,
layout: &ProjectLayout,
provider: Option<&str>,
progress: &dyn Fn(&str),
) -> Result<DriftReport> {
let descs = gather(store, hierarchy, &cfg.drift, &cfg.language);
let comparable: Vec<&EntityDescriptions> =
descs.iter().filter(|d| d.snippets.len() >= 2).collect();
if comparable.is_empty() {
return Err(Error::Store(
"drift scan: no entity has two or more retrievable descriptions to compare — \
populate the entity books and make sure the vector index is built"
.into(),
));
}
let language = if cfg.language.trim().is_empty() {
"English".to_string()
} else {
cfg.language.clone()
};
let ai = AiClient::from_config(&cfg.llm)?;
let (model, _env) = ai.resolve_provider(&cfg.llm, provider)?;
progress(&format!(
"drift scan · language: {language} · model: {model} · {} entit{} to check",
comparable.len(),
if comparable.len() == 1 { "y" } else { "ies" }
));
let mut conflicts = Vec::new();
for (i, d) in comparable.iter().enumerate() {
progress(&format!("drift [{}/{}] {}", i + 1, comparable.len(), d.entity));
let prompt = build_drift_prompt(&language, d);
let raw = run_blocking(&ai, model, DRIFT_SYSTEM_PROMPT, &prompt)?;
let pairs = parse_drift_pairs(&raw, d.snippets.len());
conflicts.extend(resolve_conflicts(&d.entity, d.kind, &d.snippets, &pairs));
}
let report = DriftReport {
version: env!("CARGO_PKG_VERSION").to_string(),
content_hash: DriftReport::compute_hash(&descs),
conflicts,
descriptions: descs,
manuscript_fingerprint: crate::world_report::manuscript_fingerprint(hierarchy),
};
report
.save(&layout.root)
.map_err(|e| Error::Store(format!("drift save: {e}")))?;
Ok(report)
}
fn build_drift_prompt(language: &str, d: &EntityDescriptions) -> String {
let mut body = format!(
"Language: {language}.\nEntity: {} ({}).\nDescriptions, in chapter order:\n",
d.entity,
d.kind.label()
);
for (i, s) in d.snippets.iter().enumerate() {
body.push_str(&format!("[{}] (ch. {}) {}\n", i + 1, s.chapter, s.text));
}
body
}
fn run_blocking(ai: &AiClient, model: &str, system: &str, prompt: &str) -> Result<String> {
crate::ai::stream::collect_blocking(
ai.client.clone(),
model.to_string(),
Some(system.to_string()),
prompt.to_string(),
)
.map_err(|e| Error::Store(format!("inference error: {e}")))
}