use std::collections::HashSet;
use std::path::PathBuf;
use anyhow::Result;
use parking_lot::{Mutex as ParkingMutex, RwLock};
pub type FileChangeCallback = Box<dyn Fn(&str, FileChange) + Send + Sync>;
use crate::backlinks::{Backlink, BacklinkIndex, LinkGraph};
use crate::chat::{delete_chat_msg, move_from_chat, read_chat_msgs, rename_chat_msg};
use crate::checklist::{
add_checklist_item, checklist_items, complete_checklist_item, incomplete_checklist_items,
remove_checklist_item, remove_completed_checklist_items,
};
use crate::fs::VirtualFs;
use crate::habits::{habits, last_week_habits, write_habits};
use crate::html::markdown_to_html;
use crate::i18n::emoji_for;
use crate::journal::{add_emoji as journal_add_emoji, add_record as journal_add_record};
use crate::parser::{extract_headings, similar};
use crate::plugins::world_clock_for_names;
use crate::stats::{done_today, today_report};
use crate::types::{FileEntry, Habits, KnowledgeConfig, CHAT_FILENAME, DIR_USER_ROOT};
use crate::worker::{move_due_tasks, remove_completed_items};
use crate::{today_chat_header, today_journal_filename};
#[derive(Debug, Clone)]
pub enum FileChange {
Created(String),
Updated(String),
Deleted(String),
Moved {
old: String,
new: String,
},
}
#[derive(Debug, Clone)]
pub struct NoteHit {
pub path: String,
pub name: String,
pub snippet: String,
pub backlink_count: usize,
pub name_similarity: i32,
}
pub struct KnowledgeBase {
fs: RwLock<VirtualFs>,
backlinks: RwLock<BacklinkIndex>,
agent_writes: ParkingMutex<HashSet<String>>,
on_change: RwLock<Vec<FileChangeCallback>>,
}
impl std::fmt::Debug for KnowledgeBase {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KnowledgeBase")
.field("root", &self.fs.read().root())
.finish()
}
}
impl KnowledgeBase {
pub fn new(root: PathBuf) -> Result<Self> {
let fs = VirtualFs::new(root)?;
Ok(Self {
fs: RwLock::new(fs),
backlinks: RwLock::new(BacklinkIndex::new()),
agent_writes: ParkingMutex::new(HashSet::new()),
on_change: RwLock::new(Vec::new()),
})
}
pub fn for_space(space_dir: &std::path::Path) -> Result<Self> {
Self::new(space_dir.join("knowledge"))
}
pub fn root(&self) -> PathBuf {
self.fs.read().root().to_path_buf()
}
pub fn on_file_change<F>(&self, f: F)
where
F: Fn(&str, FileChange) + Send + Sync + 'static,
{
self.on_change.write().push(Box::new(f));
}
fn notify_change(&self, path: &str, change: FileChange) {
for cb in self.on_change.read().iter() {
cb(path, change.clone());
}
}
pub fn note_read(&self, path: &str) -> Result<Option<String>> {
let fs = self.fs.read();
match fs.read_path(path) {
Ok(content) => Ok(Some(content)),
Err(_) => Ok(None),
}
}
pub fn note_write(&self, path: &str, content: &str) -> Result<()> {
let fs = self.fs.read();
let is_new = fs.read_path(path).is_err();
fs.write_path(path, content)?;
{
let mut backlinks = self.backlinks.write();
backlinks.remove_file(path);
backlinks.index_file(path, content);
}
self.notify_change(
path,
if is_new {
FileChange::Created(path.to_string())
} else {
FileChange::Updated(path.to_string())
},
);
Ok(())
}
pub fn note_delete(&self, path: &str) -> Result<()> {
self.fs.read().delete_path(path)?;
self.backlinks.write().remove_file(path);
self.notify_change(path, FileChange::Deleted(path.to_string()));
Ok(())
}
pub fn note_restore(&self, path: &str, content: &str) -> Result<()> {
self.fs.read().write_path(path, content)?;
let mut backlinks = self.backlinks.write();
backlinks.remove_file(path);
backlinks.index_file(path, content);
Ok(())
}
pub fn note_move(&self, old_path: &str, new_path: &str) -> Result<()> {
self.fs.read().rename_path(old_path, new_path)?;
self.backlinks.write().remove_file(old_path);
if let Some(content) = self.note_read(new_path)? {
self.backlinks.write().index_file(new_path, &content);
}
self.notify_change(
old_path,
FileChange::Moved {
old: old_path.to_string(),
new: new_path.to_string(),
},
);
Ok(())
}
pub fn note_tree(&self, dir: &str) -> Result<Vec<FileEntry>> {
let fs = self.fs.read();
let dir = if dir.is_empty() || dir == "/" {
DIR_USER_ROOT
} else {
dir
};
Ok(fs.files_and_dirs(dir)?)
}
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<NoteHit>> {
let fs = self.fs.read();
let files = fs.search_files_by_name(query)?;
let hits: Vec<NoteHit> = files
.into_iter()
.take(limit)
.map(|f| {
let path = if f.parent_dir == DIR_USER_ROOT || f.parent_dir == "/" {
f.name.clone()
} else {
format!("{}/{}", f.parent_dir, f.name)
};
let name_sim = similar(&f.display_name, query) as i32;
let bl_count = self.backlinks.read().backlink_count(&path);
NoteHit {
path,
name: f.display_name,
snippet: String::new(),
backlink_count: bl_count,
name_similarity: name_sim,
}
})
.collect();
Ok(hits)
}
pub fn backlinks_for(&self, path: &str) -> Vec<Backlink> {
self.backlinks.read().backlinks_for(path)
}
pub fn link_graph(&self) -> LinkGraph {
self.backlinks.read().link_graph()
}
pub fn index_all(&self) -> Result<usize> {
let fs = self.fs.read();
let entries = fs.files_and_dirs(DIR_USER_ROOT)?;
let mut count = 0;
for entry in &entries {
if entry.is_dir {
let sub = fs.files_and_dirs(&entry.name)?;
for sub_entry in &sub {
if !sub_entry.is_dir && sub_entry.name.ends_with(".md") {
let path = format!("{}/{}", entry.name, sub_entry.name);
if let Ok(content) = fs.read_path(&path) {
self.backlinks.write().index_file(&path, &content);
count += 1;
}
}
}
} else if entry.name.ends_with(".md") {
if let Ok(content) = fs.read_path(&entry.name) {
self.backlinks.write().index_file(&entry.name, &content);
count += 1;
}
}
}
tracing::info!(files = count, "Knowledge base indexed");
Ok(count)
}
pub fn chat_append(&self, message: &str) -> Result<()> {
let header = today_chat_header();
let timestamp = chrono::Local::now().format("`15:04`").to_string();
let entry = format!("{timestamp} {message}");
let mut content = self.note_read(CHAT_FILENAME)?.unwrap_or_default();
if !content.contains(&header) {
if !content.trim_end().ends_with('\n') {
content.push('\n');
}
content.push_str(&header);
content.push('\n');
}
content.push_str(&entry);
content.push('\n');
self.note_write(CHAT_FILENAME, &content)?;
Ok(())
}
pub fn chat_messages(&self) -> Result<Vec<String>> {
let content = self.note_read(CHAT_FILENAME)?.unwrap_or_default();
Ok(read_chat_msgs(&content))
}
pub fn chat_delete(&self, msg_hash: &str) -> Result<bool> {
let content = self.note_read(CHAT_FILENAME)?.unwrap_or_default();
match delete_chat_msg(&content, msg_hash) {
Ok(new_content) => {
self.note_write(CHAT_FILENAME, &new_content)?;
Ok(true)
}
Err(_) => Ok(false),
}
}
pub fn chat_rename(&self, msg_hash: &str, new_body: &str) -> Result<bool> {
let content = self.note_read(CHAT_FILENAME)?.unwrap_or_default();
match rename_chat_msg(&content, msg_hash, new_body) {
Ok(new_content) => {
self.note_write(CHAT_FILENAME, &new_content)?;
Ok(true)
}
Err(_) => Ok(false),
}
}
pub fn chat_move_to(&self, msg_hash: &str, target_path: &str) -> Result<bool> {
let chat_content = self.note_read(CHAT_FILENAME)?.unwrap_or_default();
let target_content = self.note_read(target_path)?.unwrap_or_default();
let (new_chat, new_target) = move_from_chat(&chat_content, msg_hash, &target_content);
if new_chat != chat_content {
self.note_write(CHAT_FILENAME, &new_chat)?;
self.note_write(target_path, &new_target)?;
Ok(true)
} else {
Ok(false)
}
}
pub fn journal_add_record(&self, record: &str) -> Result<()> {
let fs = self.fs.read();
let tz = chrono::Local::now().offset().to_owned();
journal_add_record(&fs, record, tz)?;
Ok(())
}
pub fn journal_add_emoji(&self, emoji: &str) -> Result<()> {
let fs = self.fs.read();
let tz = chrono::Local::now().offset().to_owned();
journal_add_emoji(&fs, emoji, tz)?;
Ok(())
}
pub fn journal_today_path(&self) -> String {
let tz = chrono::Local::now().offset().to_owned();
today_journal_filename(tz)
}
pub fn habits(&self, year: i32) -> Result<Habits> {
let fs = self.fs.read();
Ok(habits(&fs, year)?)
}
pub fn habits_last_week(&self) -> Result<Habits> {
let fs = self.fs.read();
let tz = chrono::Local::now().offset().to_owned();
Ok(last_week_habits(&fs, tz)?)
}
pub fn habits_write(&self, year: i32, habits: &Habits) -> Result<()> {
let fs = self.fs.read();
write_habits(&fs, year, habits)?;
Ok(())
}
pub fn config(&self) -> Result<KnowledgeConfig> {
let fs = self.fs.read();
match fs.read_path("config.json") {
Ok(content) => Ok(serde_json::from_str(&content).unwrap_or_default()),
Err(_) => Ok(KnowledgeConfig::default()),
}
}
pub fn set_config(&self, config: &KnowledgeConfig) -> Result<()> {
let json = serde_json::to_string_pretty(config)?;
self.note_write("config.json", &json)?;
Ok(())
}
pub fn checklist_items(
&self,
path: &str,
) -> Result<(Vec<String>, std::collections::HashMap<String, bool>)> {
let content = self.note_read(path)?.unwrap_or_default();
Ok(checklist_items(&content))
}
pub fn checklist_incomplete(&self, path: &str) -> Result<Vec<String>> {
let content = self.note_read(path)?.unwrap_or_default();
Ok(incomplete_checklist_items(&content))
}
pub fn checklist_add(&self, path: &str, item: &str, checked: bool) -> Result<()> {
let content = self.note_read(path)?.unwrap_or_default();
let updated = add_checklist_item(&content, item, checked);
self.note_write(path, &updated)
}
pub fn checklist_complete(&self, path: &str, item_hash: &str) -> Result<bool> {
let content = self.note_read(path)?.unwrap_or_default();
let (new_content, found) = complete_checklist_item(&content, item_hash);
if !found.is_empty() {
self.note_write(path, &new_content)?;
Ok(true)
} else {
Ok(false)
}
}
pub fn checklist_remove(&self, path: &str, item_or_hash: &str) -> Result<bool> {
let content = self.note_read(path)?.unwrap_or_default();
let (new_content, removed) = remove_checklist_item(&content, item_or_hash);
if !removed.is_empty() {
self.note_write(path, &new_content)?;
Ok(true)
} else {
Ok(false)
}
}
pub fn checklist_remove_completed(&self, path: &str) -> Result<(String, String)> {
let content = self.note_read(path)?.unwrap_or_default();
let (kept, removed) = remove_completed_checklist_items(&content);
if !removed.is_empty() {
self.note_write(path, &kept)?;
}
Ok((kept, removed))
}
pub fn run_nightly_cleanup(&self) -> Result<crate::worker::NightlyReport> {
let fs = self.fs.read();
let config = self.config()?;
Ok(remove_completed_items(&fs, &config)?)
}
pub fn run_scheduled_tasks(&self) -> Result<Vec<String>> {
let fs = self.fs.read();
let mut config = self.config()?;
let moved = move_due_tasks(&fs, &mut config)?;
if !moved.is_empty() {
self.set_config(&config)?;
}
Ok(moved)
}
pub fn today_report(&self) -> Result<crate::stats::TodayReport> {
let fs = self.fs.read();
Ok(today_report(&fs)?)
}
pub fn done_today(&self) -> Result<Vec<FileEntry>> {
let fs = self.fs.read();
Ok(done_today(&fs)?)
}
pub fn markdown_to_html(&self, md: &str) -> String {
markdown_to_html(md)
}
pub fn auto_emoji(&self, text: &str) -> String {
emoji_for(text)
}
pub fn world_clock(&self, timezone_names: &[&str]) -> Vec<crate::plugins::TimezoneEntry> {
world_clock_for_names(timezone_names)
}
pub fn mark_agent_write(&self, path: &str) {
self.agent_writes.lock().insert(path.to_string());
}
pub fn is_agent_write(&self, path: &str) -> bool {
self.agent_writes.lock().contains(path)
}
pub fn clear_agent_write(&self, path: &str) {
self.agent_writes.lock().remove(path);
}
pub fn extract_text_imgs_links(&self, text: &str) -> crate::tgtxt::ExtractResult {
crate::tgtxt::extract_text_imgs_links(text)
}
pub fn extract_headings(&self, content: &str) -> Vec<String> {
extract_headings(content).into_iter().take(5).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_test_kb() -> KnowledgeBase {
let dir = std::env::temp_dir().join(format!("test-kb-{}", uuid::Uuid::new_v4()));
KnowledgeBase::new(dir.join("kb")).expect("test knowledge base")
}
#[test]
fn test_note_write_and_read() {
let kb = make_test_kb();
kb.note_write("brain/Rust.md", "# Rust\n\nHello world")
.unwrap();
let content = kb.note_read("brain/Rust.md").unwrap();
assert_eq!(content, Some("# Rust\n\nHello world".to_string()));
}
#[test]
fn test_note_read_missing() {
let kb = make_test_kb();
assert_eq!(kb.note_read("nonexistent.md").unwrap(), None);
}
#[test]
fn test_note_delete() {
let kb = make_test_kb();
kb.note_write("del.md", "to delete").unwrap();
kb.note_delete("del.md").unwrap();
assert_eq!(kb.note_read("del.md").unwrap(), None);
}
#[test]
fn test_note_move() {
let kb = make_test_kb();
kb.note_write("old.md", "content").unwrap();
kb.note_move("old.md", "new.md").unwrap();
assert_eq!(kb.note_read("old.md").unwrap(), None);
assert_eq!(kb.note_read("new.md").unwrap(), Some("content".to_string()));
}
#[test]
fn test_backlinks() {
let kb = make_test_kb();
kb.note_write("brain/Rust.md", "See [Ownership](brain/Ownership.md)")
.unwrap();
let bl = kb.backlinks_for("brain/Ownership.md");
assert_eq!(bl.len(), 1);
assert_eq!(bl[0].source_path, "brain/Rust.md");
}
#[test]
fn test_note_tree() {
let kb = make_test_kb();
kb.note_write("brain/Rust.md", "Rust").unwrap();
let entries = kb.note_tree("brain").unwrap();
assert!(!entries.is_empty());
}
#[test]
fn test_search_by_name() {
let kb = make_test_kb();
kb.note_write("brain/Rust.md", "Rust content").unwrap();
let hits = kb.search("Rust", 10).unwrap();
assert!(!hits.is_empty());
}
#[test]
fn test_link_graph() {
let kb = make_test_kb();
kb.note_write("a.md", "[b](b.md)").unwrap();
let graph = kb.link_graph();
assert!(!graph.edges.is_empty());
}
#[test]
fn test_agent_write_tracking() {
let kb = make_test_kb();
assert!(!kb.is_agent_write("test.md"));
kb.mark_agent_write("test.md");
assert!(kb.is_agent_write("test.md"));
kb.clear_agent_write("test.md");
assert!(!kb.is_agent_write("test.md"));
}
#[test]
fn test_index_all() {
let kb = make_test_kb();
kb.note_write("brain/Rust.md", "Rust [Go](brain/Go.md)")
.unwrap();
kb.note_write("brain/Go.md", "Go language").unwrap();
kb.note_write("index.md", "Welcome").unwrap();
let count = kb.index_all().unwrap();
assert_eq!(count, 3);
let bl = kb.backlinks_for("brain/Go.md");
assert_eq!(bl.len(), 1);
}
#[test]
fn test_on_file_change_callback() {
let kb = make_test_kb();
let _called = std::sync::atomic::AtomicBool::new(false);
let path_clone: std::sync::Arc<std::sync::atomic::AtomicBool> =
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let flag = path_clone.clone();
kb.on_file_change(move |path, change| {
let _ = path;
let _ = change;
flag.store(true, std::sync::atomic::Ordering::SeqCst);
});
kb.note_write("test.md", "hello").unwrap();
assert!(path_clone.load(std::sync::atomic::Ordering::SeqCst));
}
#[test]
fn test_chat_append() {
let kb = make_test_kb();
kb.chat_append("Test message").unwrap();
let messages = kb.chat_messages().unwrap();
assert!(!messages.is_empty());
}
#[test]
fn test_config() {
let kb = make_test_kb();
let cfg = kb.config().unwrap();
let cfg2 = kb.config().unwrap();
assert_eq!(cfg.language, cfg2.language);
}
#[test]
fn test_markdown_to_html() {
let kb = make_test_kb();
let html = kb.markdown_to_html("# Hello\n\n**world**");
assert!(html.contains("Hello"), "HTML should contain Hello: {html}");
assert!(html.contains("world"), "HTML should contain world: {html}");
}
#[test]
fn test_auto_emoji() {
let kb = make_test_kb();
let emoji = kb.auto_emoji("cooking pasta");
assert!(!emoji.is_empty());
}
#[test]
fn test_extract_headings() {
let kb = make_test_kb();
let headings = kb.extract_headings("# Title\n\n## Section\n\n### Subsection");
assert!(headings.len() >= 2);
}
}