use crate::{storage::{self, KnowledgeBase}, text, types::*, Error, Result};
use rusqlite::{params, Connection, OptionalExtension};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::collections::BTreeMap;
use std::path::PathBuf;
fn default_chunk_chars() -> usize { 220 }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NoteFileInput {
#[serde(flatten)] pub record: RecordInput,
pub path: PathBuf,
#[serde(default = "default_chunk_chars")] pub chunk_chars: usize,
}
impl NoteFileInput {
pub fn new(path: impl Into<PathBuf>) -> Self {
Self { record: RecordInput::default(), path: path.into(), chunk_chars: default_chunk_chars() }
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Note {
#[serde(flatten)] pub header: RecordHeader,
pub source: String,
#[serde(default)] pub title: String,
pub chunk_chars: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Chunk {
#[serde(flatten)] pub header: RecordHeader,
pub note_id: i64, pub ordinal: usize, pub offset: usize, pub limit: usize,
#[serde(default)] pub content: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TextChunk { pub ordinal: usize, pub offset: usize, pub limit: usize, pub content: String }
pub fn chunk_text(content: &str, target: usize) -> Result<Vec<TextChunk>> {
if !(16..=100_000).contains(&target) { return Err(Error::Validation("chunk_chars must be between 16 and 100000".into())); }
let lines: Vec<&str> = content.lines().collect();
let mut blocks: Vec<(usize, usize, String, bool)> = Vec::new();
let mut i = 0;
while i < lines.len() {
if lines[i].trim().is_empty() { i += 1; continue; }
let start = i;
let trimmed = lines[i].trim_start();
let fence_char = trimmed.chars().next().filter(|c| *c == '`' || *c == '~');
let fence_len = fence_char.map(|c| trimmed.chars().take_while(|x| *x == c).count()).unwrap_or(0);
let is_fence = fence_len >= 3;
let is_table = lines[i].contains('|') && i + 1 < lines.len() && {
let next = lines[i + 1].trim();
next.contains('-') && next.contains('|') && next.chars().all(|c| matches!(c, '-' | ':' | '|' | ' ' | '\t'))
};
i += 1;
if is_fence {
while i < lines.len() {
let line = lines[i].trim();
i += 1;
if line.chars().take_while(|c| Some(*c) == fence_char).count() >= fence_len
&& line.chars().all(|c| Some(c) == fence_char || c.is_whitespace()) { break; }
}
} else if is_table {
while i < lines.len() && lines[i].contains('|') && !lines[i].trim().is_empty() { i += 1; }
} else {
while i < lines.len() && !lines[i].trim().is_empty() {
let line = lines[i].trim_start();
if line.starts_with("```") || line.starts_with("~~~") || line.starts_with('#') { break; }
if i + 1 < lines.len() && lines[i].contains('|') && lines[i + 1].contains("---") { break; }
i += 1;
}
}
blocks.push((start, i, lines[start..i].join("\n"), is_fence || is_table));
}
let mut chunks = Vec::new();
for (start, end, body, atomic) in blocks {
if atomic || body.chars().count() <= target {
chunks.push(TextChunk { ordinal: 0, offset: start + 1, limit: end - start, content: body });
continue;
}
let mut current = String::new();
let mut count = 0;
let mut first_line = start + 1;
let mut last_line = first_line;
for (line_no, line) in lines.iter().enumerate().take(end).skip(start) {
if !current.is_empty() {
if count + 1 >= target {
chunks.push(TextChunk { ordinal: 0, offset: first_line, limit: last_line - first_line + 1, content: std::mem::take(&mut current) });
count = 0;
} else { current.push('\n'); count += 1; }
}
for ch in line.chars() {
if count == target {
chunks.push(TextChunk { ordinal: 0, offset: first_line, limit: last_line - first_line + 1, content: std::mem::take(&mut current) });
count = 0;
}
if current.is_empty() { first_line = line_no + 1; }
current.push(ch); count += 1; last_line = line_no + 1;
}
}
if !current.is_empty() { chunks.push(TextChunk { ordinal: 0, offset: first_line, limit: last_line - first_line + 1, content: current }); }
}
for (ordinal, chunk) in chunks.iter_mut().enumerate() { chunk.ordinal = ordinal; }
Ok(chunks)
}
pub(crate) struct PreparedNote {
pub record: RecordInput,
pub namespace_id: i64, pub scope_id: i64, pub source: String,
pub given: String, pub title: String, pub split: Vec<TextChunk>,
pub chunk_chars: usize,
}
pub(crate) fn prepare_note(conn: &Connection, input: &NoteFileInput, split: Vec<TextChunk>) -> Result<PreparedNote> {
let given = input.path.to_string_lossy().into_owned();
storage::validate_identity("source", &given)?;
let namespace_id = storage::term_id(conn, &input.record.namespace)?;
let scope_id = storage::term_id(conn, &input.record.scope)?;
let root = storage::namespace_root(conn, namespace_id)?.ok_or_else(|| Error::Validation(
format!("namespace {} has no registered domain root; call notes.set_root before upserting notes", input.record.namespace)))?;
let source = relative_note_path(&root, &given)?;
let title = input.path.file_stem().map(|stem| stem.to_string_lossy().into_owned()).unwrap_or_default();
let (dirs, split_stem) = storage::split_note_path(&source);
let stem = if title.is_empty() { split_stem } else { title.clone() };
let mut record = input.record.clone();
record.id = None;
let mut tags = record.tags.clone();
tags.extend(dirs.iter().cloned());
if !stem.is_empty() { tags.push(stem.clone()); }
record.tags = tags.clone();
Ok(PreparedNote { record, namespace_id, scope_id, source, given, title, split, chunk_chars: input.chunk_chars })
}
pub(crate) fn add_note(conn: &Connection, note: &PreparedNote) -> Result<(Note, Vec<crate::index::IndexDocument>)> {
let (header, _) = storage::put_record(conn, RecordKind::Note, ¬e.record,
&json!({"chunk_chars": note.chunk_chars}), "")?;
conn.execute("INSERT INTO notes(record_id,namespace_id,scope_id,path,name) VALUES (?1,?2,?3,?4,?5)
ON CONFLICT(record_id) DO UPDATE SET namespace_id=excluded.namespace_id,scope_id=excluded.scope_id,path=excluded.path,name=excluded.name",
params![header.id, note.namespace_id, note.scope_id, note.source, note.title])?;
let mut documents = Vec::new();
for chunk in ¬e.split {
let content_digest = text::digest(&chunk.content);
let payload = json!({"note_id":header.id,"ordinal":chunk.ordinal,"offset":chunk.offset,"limit":chunk.limit});
let chunk_input = RecordInput { id: None, namespace: header.namespace.clone(), scope: header.scope.clone(),
tags: header.tags.clone(), evidence: vec![], metadata: header.metadata.clone(),
created_at_us: Some(header.created_at_us), updated_at_us: Some(header.updated_at_us), expected_revision: None };
let (chunk_header, document) = storage::put_record(conn, RecordKind::Chunk, &chunk_input, &payload, &chunk.content)?;
conn.execute("INSERT INTO chunks(record_id,note_id,ordinal,\"offset\",\"limit\",fingerprint) VALUES (?1,?2,?3,?4,?5,?6)",
params![chunk_header.id, header.id, chunk.ordinal as i64, chunk.offset as i64, chunk.limit as i64, content_digest])?;
documents.push(document);
}
Ok((Note { header, source: note.given.clone(), title: note.title.clone(), chunk_chars: note.chunk_chars }, documents))
}
fn relative_note_path(root: &str, given: &str) -> Result<String> {
let root = root.replace('\\', "/");
let root = root.trim_end_matches('/');
let relative = given.replace('\\', "/");
let Some(relative) = relative.strip_prefix(root).and_then(|rest| rest.strip_prefix('/')) else {
return Err(Error::Validation(format!("note path {given} is outside the domain root {root}")));
};
if relative.is_empty() { return Err(Error::Validation("note path must name a file below the domain root".into())); }
Ok(relative.to_string())
}
#[derive(Clone)]
pub struct NoteStore(pub(crate) KnowledgeBase);
impl NoteStore {
pub fn upsert_file(&self, input: NoteFileInput) -> Result<WriteReceipt<Note>> {
let mut receipts = self.upsert_files(&[input])?;
let value = receipts.value.pop().ok_or_else(|| Error::Validation("empty upsert batch".into()))?;
Ok(WriteReceipt { value, revision: receipts.revision })
}
pub fn upsert_files(&self, inputs: &[NoteFileInput]) -> Result<WriteReceipt<Vec<Note>>> {
let split = inputs.iter().map(|input| -> Result<Vec<TextChunk>> {
let content = std::fs::read_to_string(&input.path)?;
chunk_text(&content, input.chunk_chars)
}).collect::<Result<Vec<_>>>()?;
let mut notes: Vec<Note> = Vec::new();
let (stale, documents) = self.0.with_writer_lock(|writer| {
let mut documents: Vec<crate::index::IndexDocument> = Vec::new();
let mut stale: Vec<i64> = Vec::new();
for (input, split) in inputs.iter().zip(split) {
let tx = writer.conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let prepared = prepare_note(&tx, input, split)?;
let by_path: Option<i64> = tx.query_row("SELECT record_id FROM notes WHERE namespace_id=?1 AND scope_id=?2 AND path=?3",
params![prepared.namespace_id, prepared.scope_id, prepared.source], |r| r.get(0)).optional()?;
let mut removed: Vec<i64> = Vec::new();
if let Some(id) = by_path {
let mut stmt = tx.prepare("SELECT record_id FROM chunks WHERE note_id=?1")?;
for row in stmt.query_map([id], |r| r.get::<_, i64>(0))? { removed.push(row?); }
removed.push(id);
storage::mark_records(&tx, &removed, storage::MARK_DELETING)?;
}
tx.commit()?;
if !removed.is_empty() {
let tx = writer.conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
for id in &removed { storage::delete_record(&tx, &RecordKey { id: *id })?; }
tx.commit()?;
stale.extend(removed);
}
let tx = writer.conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let (note, docs) = add_note(&tx, &prepared)?;
tx.commit()?;
documents.extend(docs);
notes.push(note);
}
Ok((stale, documents))
})?;
if !stale.is_empty() { self.0.index()?.stage_deletions(&stale)?; }
self.0.index_documents(&documents)?;
let revision = storage::current_revision(self.0.read()?.conn())?;
Ok(WriteReceipt { value: notes, revision })
}
pub fn set_root(&self, namespace: &str, root: &str) -> Result<()> {
storage::validate_identity("namespace", namespace)?;
if !std::path::Path::new(root).is_dir() {
return Err(Error::Validation(format!("domain root {root} is not an existing directory")));
}
let root = root.replace('\\', "/");
let root = root.trim_end_matches('/').to_string();
self.0.write(|writer| {
let namespace_id = storage::term_id(&writer.conn, namespace)?;
writer.conn.execute("INSERT INTO namespace_roots(namespace_id,root) VALUES (?1,?2)
ON CONFLICT(namespace_id) DO UPDATE SET root=excluded.root", params![namespace_id, root])?;
Ok(())
})
}
pub fn unset_root(&self, namespace: &str) -> Result<bool> {
storage::validate_identity("namespace", namespace)?;
self.0.write(|writer| {
let namespace_id = storage::term_id(&writer.conn, namespace)?;
Ok(writer.conn.execute("DELETE FROM namespace_roots WHERE namespace_id=?1", [namespace_id])? > 0)
})
}
pub fn root(&self, namespace: &str) -> Result<Option<String>> {
let state = self.0.read()?;
let conn = state.conn();
let namespace_id: Option<i64> = conn.query_row("SELECT id FROM strings WHERE text=?1",
[text::normalized_tag(namespace)], |r| r.get(0)).optional()?;
let Some(namespace_id) = namespace_id else { return Ok(None) };
storage::namespace_root(conn, namespace_id)
}
pub fn get(&self, id: i64, filter: &ReadFilter) -> Result<Note> {
storage::get(self.0.read()?.conn(), &RecordKey { id }, filter)
}
pub fn get_many(&self, ids: &[i64], filter: &ReadFilter) -> Result<BTreeMap<i64, Note>> {
storage::load_many(self.0.read()?.conn(), ids, filter)
}
pub fn list(&self, page: &PageRequest) -> Result<Page<Note>> { storage::list(self.0.read()?.conn(), RecordKind::Note, page) }
pub fn get_chunk(&self, id: i64, filter: &ReadFilter) -> Result<Chunk> {
let state = self.0.read()?;
let conn = state.conn();
let mut chunk: Chunk = storage::get(conn, &RecordKey { id }, filter)?;
chunk.content = self.0.index()?.bodies(&[id])?.remove(&id).unwrap_or_default();
Ok(chunk)
}
pub fn chunks(&self, note_id: i64, filter: &ReadFilter) -> Result<Vec<Chunk>> {
let state = self.0.read()?;
let conn = state.conn();
let _: Note = storage::get(conn, &RecordKey { id: note_id }, filter)?;
let ids: Vec<i64> = {
let mut stmt = conn.prepare("SELECT record_id FROM chunks WHERE note_id=?1 ORDER BY ordinal")?;
let rows = stmt.query_map([note_id], |r| r.get::<_, i64>(0))?;
rows.collect::<std::result::Result<Vec<_>, _>>()?
};
let mut loaded: BTreeMap<i64, Chunk> = storage::load_many(conn, &ids, filter)?;
let mut chunks: Vec<Chunk> = ids.iter().filter_map(|id| loaded.remove(id)).collect();
let ids: Vec<i64> = chunks.iter().map(|chunk| chunk.header.id).collect();
let mut bodies = self.0.index()?.bodies(&ids)?;
for chunk in &mut chunks { chunk.content = bodies.remove(&chunk.header.id).unwrap_or_default(); }
Ok(chunks)
}
pub fn delete(&self, ids: &[i64], filter: &ReadFilter) -> Result<WriteReceipt<usize>> {
let removed = self.0.delete_flow(ids, filter, RecordKind::Note, |tx, id| {
let mut out = Vec::new();
let mut stmt = tx.prepare("SELECT record_id FROM chunks WHERE note_id=?1")?;
for row in stmt.query_map([id], |r| r.get::<_, i64>(0))? { out.push(row?); }
Ok(out)
})?;
let revision = storage::current_revision(self.0.read()?.conn())?;
Ok(WriteReceipt { value: removed, revision })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn split_note_path_separates_directories_from_the_file_name() {
assert_eq!(storage::split_note_path("notes/characters/overview.md"), (vec!["notes".to_string(), "characters".to_string()], "overview".to_string()));
assert_eq!(storage::split_note_path("v1.2/角色.设定.md"), (vec!["v1.2".to_string()], "角色.设定".to_string()));
assert_eq!(storage::split_note_path("overview"), (Vec::new(), "overview".to_string()));
}
#[test]
fn note_paths_must_stay_inside_the_domain_root() {
assert_eq!(relative_note_path("E:/data/demo", r"E:\data\demo\notes\overview.md").unwrap(), "notes/overview.md");
assert!(relative_note_path("E:/data/demo", "E:/data/demox/overview.md").is_err());
assert!(relative_note_path("E:/data/demo", "E:/data/other/overview.md").is_err());
assert!(relative_note_path("E:/data/demo", "E:/data/demo").is_err());
}
}