Skip to main content

ninox_core/
brain.rs

1use anyhow::{Context, Result};
2use rusqlite::{params, Connection};
3use serde::{Deserialize, Serialize};
4use std::{
5    collections::HashMap,
6    fs,
7    path::{Path, PathBuf},
8    sync::Mutex,
9};
10use walkdir::WalkDir;
11
12// ---------------------------------------------------------------------------
13// Public types
14// ---------------------------------------------------------------------------
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct BrainEntry {
18    pub id: String,         // relative path from brain root (e.g. "people/alice.md")
19    pub entry_type: String, // derived from parent directory name
20    pub name: String,       // from frontmatter or filename stem
21    pub tags: Vec<String>,
22    pub repos: Vec<String>,
23    pub updated: Option<String>,
24    pub body: String,
25}
26
27#[derive(Debug, Default)]
28pub struct QueryFilters {
29    pub entry_type: Option<String>,
30    pub tag: Option<String>,
31}
32
33// ---------------------------------------------------------------------------
34// BrainIndex
35// ---------------------------------------------------------------------------
36
37pub struct BrainIndex {
38    conn: Mutex<Connection>,
39    brain_path: PathBuf,
40}
41
42impl BrainIndex {
43    pub fn open(brain_path: impl AsRef<Path>) -> Result<Self> {
44        let brain_path = brain_path.as_ref().to_path_buf();
45        fs::create_dir_all(&brain_path)
46            .with_context(|| format!("create brain dir {brain_path:?}"))?;
47        let db_path = brain_path.join(".index.db");
48        let conn = Connection::open(&db_path)
49            .with_context(|| format!("open brain db {db_path:?}"))?;
50        conn.execute_batch(
51            "PRAGMA journal_mode=WAL;
52             CREATE TABLE IF NOT EXISTS entries (
53                 id      TEXT PRIMARY KEY,
54                 type    TEXT NOT NULL,
55                 name    TEXT NOT NULL,
56                 tags    TEXT NOT NULL DEFAULT '[]',
57                 repos   TEXT NOT NULL DEFAULT '[]',
58                 updated TEXT,
59                 body    TEXT NOT NULL DEFAULT ''
60             );
61             CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts
62                 USING fts5(name, tags, body, content=entries, content_rowid=rowid);",
63        )?;
64        ensure_gitignore(&brain_path)?;
65        Ok(Self { conn: Mutex::new(conn), brain_path })
66    }
67
68    /// Walk the brain directory, parse markdown files, and repopulate the index.
69    /// Returns the number of files indexed.
70    pub fn rebuild(&self) -> Result<usize> {
71        let conn = self.conn.lock().unwrap();
72        conn.execute_batch("DELETE FROM entries; DELETE FROM entries_fts;")?;
73
74        let mut count = 0usize;
75        for entry in WalkDir::new(&self.brain_path)
76            .follow_links(false)
77            .into_iter()
78            .filter_map(|e| e.ok())
79        {
80            let path = entry.path();
81            // Skip non-files, non-.md files, and the index db itself
82            if !path.is_file() {
83                continue;
84            }
85            let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
86            if ext != "md" {
87                continue;
88            }
89
90            let rel = path
91                .strip_prefix(&self.brain_path)
92                .unwrap_or(path)
93                .to_string_lossy()
94                .to_string();
95
96            let content = match fs::read_to_string(path) {
97                Ok(c) => c,
98                Err(_) => continue,
99            };
100
101            let parsed = parse_markdown(&content);
102
103            // Derive entry_type from parent dir name (or frontmatter "type" field)
104            let parent_type = path
105                .parent()
106                .and_then(|p| {
107                    // If the parent IS brain_path itself, there's no meaningful type dir
108                    if p == self.brain_path { None } else { p.file_name() }
109                })
110                .and_then(|n| n.to_str())
111                .map(str::to_string);
112
113            let entry_type = parsed
114                .frontmatter
115                .get("type")
116                .and_then(|v| v.as_str())
117                .map(str::to_string)
118                .or(parent_type)
119                .unwrap_or_else(|| "note".to_string());
120
121            let stem = path
122                .file_stem()
123                .and_then(|s| s.to_str())
124                .unwrap_or("unknown");
125            let name = parsed
126                .frontmatter
127                .get("name")
128                .and_then(|v| v.as_str())
129                .map(str::to_string)
130                .unwrap_or_else(|| stem.to_string());
131
132            let tags: Vec<String> = parsed
133                .frontmatter
134                .get("tags")
135                .and_then(|v| v.as_sequence())
136                .cloned()
137                .unwrap_or_default();
138
139            let repos: Vec<String> = parsed
140                .frontmatter
141                .get("repos")
142                .and_then(|v| v.as_sequence())
143                .cloned()
144                .unwrap_or_default();
145
146            let updated = parsed
147                .frontmatter
148                .get("updated")
149                .and_then(|v| v.as_str())
150                .map(str::to_string);
151
152            let tags_json = serde_json::to_string(&tags)?;
153            let repos_json = serde_json::to_string(&repos)?;
154
155            conn.execute(
156                "INSERT INTO entries (id, type, name, tags, repos, updated, body)
157                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
158                params![rel, entry_type, name, tags_json, repos_json, updated, parsed.body],
159            )?;
160
161            count += 1;
162        }
163
164        // Rebuild the FTS index from the content table
165        conn.execute_batch("INSERT INTO entries_fts(entries_fts) VALUES('rebuild');")?;
166
167        Ok(count)
168    }
169
170    /// Full-text search with optional filters.
171    pub fn query(&self, text: &str, filters: QueryFilters) -> Result<Vec<BrainEntry>> {
172        let conn = self.conn.lock().unwrap();
173
174        if text.is_empty() && filters.entry_type.is_none() && filters.tag.is_none() {
175            // Return all entries when no constraints given
176            let mut stmt = conn.prepare(
177                "SELECT id, type, name, tags, repos, updated, body FROM entries ORDER BY name",
178            )?;
179            let rows = stmt.query_map([], row_to_entry)?;
180            let entries: Vec<BrainEntry> =
181                rows.collect::<rusqlite::Result<Vec<_>>>()?;
182            return Ok(entries);
183        }
184
185        if text.is_empty() {
186            // Filter-only query
187            let mut stmt = conn.prepare(
188                "SELECT id, type, name, tags, repos, updated, body FROM entries
189                 WHERE (?1 IS NULL OR type = ?1)
190                 ORDER BY name",
191            )?;
192            let rows = stmt.query_map(params![filters.entry_type.as_deref()], row_to_entry)?;
193            let mut results: Vec<BrainEntry> =
194                rows.collect::<rusqlite::Result<Vec<_>>>()?;
195            if let Some(ref tag) = filters.tag {
196                results.retain(|e| e.tags.iter().any(|t| t == tag));
197            }
198            return Ok(results);
199        }
200
201        let mut stmt = conn.prepare(
202            "SELECT e.id, e.type, e.name, e.tags, e.repos, e.updated, e.body
203             FROM entries_fts
204             JOIN entries e ON entries_fts.rowid = e.rowid
205             WHERE entries_fts MATCH ?1
206             ORDER BY rank",
207        )?;
208        let rows = stmt.query_map(params![text], row_to_entry)?;
209        let mut results: Vec<BrainEntry> =
210            rows.collect::<rusqlite::Result<Vec<_>>>()?;
211
212        // Post-filter by type and tag
213        if let Some(ref et) = filters.entry_type {
214            results.retain(|e| &e.entry_type == et);
215        }
216        if let Some(ref tag) = filters.tag {
217            results.retain(|e| e.tags.iter().any(|t| t == tag));
218        }
219
220        Ok(results)
221    }
222
223    /// Fetch a single entry by its relative path id.
224    pub fn get(&self, id: &str) -> Result<Option<BrainEntry>> {
225        let conn = self.conn.lock().unwrap();
226        let mut stmt = conn.prepare(
227            "SELECT id, type, name, tags, repos, updated, body FROM entries WHERE id = ?1",
228        )?;
229        let mut rows = stmt.query_map([id], row_to_entry)?;
230        match rows.next() {
231            None => Ok(None),
232            Some(r) => Ok(Some(r?)),
233        }
234    }
235}
236
237// ---------------------------------------------------------------------------
238// Helpers
239// ---------------------------------------------------------------------------
240
241fn row_to_entry(r: &rusqlite::Row) -> rusqlite::Result<BrainEntry> {
242    let tags_json: String = r.get(3)?;
243    let repos_json: String = r.get(4)?;
244    let tags: Vec<String> = serde_json::from_str(&tags_json).unwrap_or_default();
245    let repos: Vec<String> = serde_json::from_str(&repos_json).unwrap_or_default();
246    Ok(BrainEntry {
247        id: r.get(0)?,
248        entry_type: r.get(1)?,
249        name: r.get(2)?,
250        tags,
251        repos,
252        updated: r.get(5)?,
253        body: r.get(6)?,
254    })
255}
256
257/// Ensure `.index.db` is in the brain directory's `.gitignore`.
258fn ensure_gitignore(brain_path: &Path) -> Result<()> {
259    let gi = brain_path.join(".gitignore");
260    let entry = ".index.db\n";
261    if gi.exists() {
262        let content = fs::read_to_string(&gi)?;
263        if !content.contains(".index.db") {
264            fs::write(&gi, format!("{content}{entry}"))?;
265        }
266    } else {
267        fs::write(&gi, entry)?;
268    }
269    Ok(())
270}
271
272// ---------------------------------------------------------------------------
273// Frontmatter parsing (manual YAML split, no extra dep)
274// ---------------------------------------------------------------------------
275
276struct FmValue {
277    str_val: Option<String>,
278    seq_val: Option<Vec<String>>,
279}
280
281impl FmValue {
282    fn str(s: &str) -> Self {
283        Self { str_val: Some(s.to_string()), seq_val: None }
284    }
285
286    fn seq(v: Vec<String>) -> Self {
287        Self { str_val: None, seq_val: Some(v) }
288    }
289
290    fn as_str(&self) -> Option<&str> {
291        self.str_val.as_deref()
292    }
293
294    fn as_sequence(&self) -> Option<&Vec<String>> {
295        self.seq_val.as_ref()
296    }
297}
298
299struct Frontmatter(HashMap<String, FmValue>);
300
301impl Frontmatter {
302    fn get(&self, key: &str) -> Option<&FmValue> {
303        self.0.get(key)
304    }
305}
306
307struct ParsedMd {
308    frontmatter: Frontmatter,
309    body: String,
310}
311
312fn parse_markdown(content: &str) -> ParsedMd {
313    if !content.starts_with("---") {
314        return ParsedMd {
315            frontmatter: Frontmatter(HashMap::new()),
316            body: content.to_string(),
317        };
318    }
319    let rest = &content[3..];
320    let end = rest.find("\n---").or_else(|| rest.find("\r\n---"));
321    let (fm_text, body) = match end {
322        None => ("", content),
323        Some(pos) => {
324            let after = &rest[pos + 4..]; // skip "\n---"
325            // skip optional trailing newline
326            let body = after.trim_start_matches('\n').trim_start_matches('\r');
327            (&rest[..pos], body)
328        }
329    };
330    let fm = parse_frontmatter(fm_text);
331    ParsedMd { frontmatter: fm, body: body.to_string() }
332}
333
334fn parse_frontmatter(text: &str) -> Frontmatter {
335    let mut map: HashMap<String, FmValue> = HashMap::new();
336    let mut lines = text.lines().peekable();
337    while let Some(line) = lines.next() {
338        let line = line.trim();
339        if line.is_empty() {
340            continue;
341        }
342        if let Some((key, val)) = line.split_once(':') {
343            let key = key.trim().to_string();
344            let val = val.trim();
345            if val.is_empty() {
346                // Possibly a sequence starting on the next lines
347                let mut seq = Vec::new();
348                while let Some(next) = lines.peek() {
349                    let t = next.trim();
350                    if let Some(stripped) = t.strip_prefix("- ") {
351                        seq.push(stripped.trim().to_string());
352                        lines.next();
353                    } else {
354                        break;
355                    }
356                }
357                if !seq.is_empty() {
358                    map.insert(key, FmValue::seq(seq));
359                }
360            } else if val.starts_with('[') && val.ends_with(']') {
361                // Inline sequence: [a, b, c]
362                let inner = &val[1..val.len() - 1];
363                let seq: Vec<String> = inner
364                    .split(',')
365                    .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_string())
366                    .filter(|s| !s.is_empty())
367                    .collect();
368                map.insert(key, FmValue::seq(seq));
369            } else {
370                map.insert(
371                    key,
372                    FmValue::str(val.trim_matches('"').trim_matches('\'')),
373                );
374            }
375        }
376    }
377    Frontmatter(map)
378}
379
380// ---------------------------------------------------------------------------
381// Tests
382// ---------------------------------------------------------------------------
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387    use tempfile::tempdir;
388
389    fn make_brain() -> (BrainIndex, tempfile::TempDir) {
390        let dir = tempdir().unwrap();
391        let brain = BrainIndex::open(dir.path()).unwrap();
392        (brain, dir)
393    }
394
395    #[test]
396    fn open_creates_schema() {
397        let (_brain, dir) = make_brain();
398        let db_path = dir.path().join(".index.db");
399        assert!(db_path.exists());
400        // Verify the gitignore was created
401        let gi = dir.path().join(".gitignore");
402        assert!(gi.exists());
403        let content = fs::read_to_string(&gi).unwrap();
404        assert!(content.contains(".index.db"));
405    }
406
407    #[test]
408    fn rebuild_indexes_files() {
409        let (brain, dir) = make_brain();
410        let people_dir = dir.path().join("people");
411        fs::create_dir_all(&people_dir).unwrap();
412        fs::write(
413            people_dir.join("alice.md"),
414            "---\nname: Alice Smith\ntags:\n- engineering\n- leadership\n---\nAlice leads the infra team.",
415        )
416        .unwrap();
417        fs::write(
418            people_dir.join("bob.md"),
419            "# Bob\n\nBob works on frontend.",
420        )
421        .unwrap();
422
423        let count = brain.rebuild().unwrap();
424        assert_eq!(count, 2);
425    }
426
427    #[test]
428    fn query_returns_matches() {
429        let (brain, dir) = make_brain();
430        let dir_path = dir.path().join("notes");
431        fs::create_dir_all(&dir_path).unwrap();
432        fs::write(
433            dir_path.join("rust-tips.md"),
434            "---\nname: Rust Tips\ntags:\n- rust\n---\nUse anyhow for error handling.",
435        )
436        .unwrap();
437        fs::write(
438            dir_path.join("python-tips.md"),
439            "---\nname: Python Tips\ntags:\n- python\n---\nUse dataclasses for data.",
440        )
441        .unwrap();
442
443        brain.rebuild().unwrap();
444
445        let results = brain.query("anyhow", QueryFilters::default()).unwrap();
446        assert_eq!(results.len(), 1);
447        assert_eq!(results[0].name, "Rust Tips");
448    }
449
450    #[test]
451    fn query_filters_by_type() {
452        let (brain, dir) = make_brain();
453        let people = dir.path().join("people");
454        let projects = dir.path().join("projects");
455        fs::create_dir_all(&people).unwrap();
456        fs::create_dir_all(&projects).unwrap();
457        fs::write(people.join("alice.md"), "Alice is a person.").unwrap();
458        fs::write(projects.join("athene.md"), "Athene is a project.").unwrap();
459
460        brain.rebuild().unwrap();
461
462        let results = brain
463            .query("", QueryFilters { entry_type: Some("people".into()), tag: None })
464            .unwrap();
465        assert_eq!(results.len(), 1);
466        assert_eq!(results[0].entry_type, "people");
467    }
468}