org_core/
org_mode.rs

1use std::{fs, io, path::PathBuf};
2
3use nucleo_matcher::pattern::{AtomKind, CaseMatching, Normalization, Pattern};
4use nucleo_matcher::{Config as NucleoConfig, Matcher};
5use orgize::Org;
6use orgize::ast::PropertyDrawer;
7use orgize::export::{Container, Event, from_fn, from_fn_with_ctx};
8use serde::{Deserialize, Serialize};
9use walkdir::WalkDir;
10
11use crate::OrgModeError;
12use crate::config::{OrgConfig, load_org_config};
13
14#[cfg(test)]
15#[path = "org_mode_tests.rs"]
16mod org_mode_tests;
17
18#[derive(Debug)]
19pub struct OrgMode {
20    config: OrgConfig,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct TreeNode {
25    pub label: String,
26    pub level: usize,
27    #[serde(skip_serializing_if = "Vec::is_empty", default)]
28    pub children: Vec<TreeNode>,
29    #[serde(skip_serializing_if = "Vec::is_empty", default)]
30    pub tags: Vec<String>,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct SearchResult {
35    pub file_path: String,
36    pub snippet: String,
37    pub score: u32,
38    #[serde(skip_serializing_if = "Vec::is_empty", default)]
39    pub tags: Vec<String>,
40}
41
42impl TreeNode {
43    pub fn new(label: String) -> Self {
44        Self {
45            label,
46            level: 0,
47            children: Vec::new(),
48            tags: Vec::new(),
49        }
50    }
51
52    pub fn new_with_level(label: String, level: usize) -> Self {
53        Self {
54            label,
55            level,
56            children: Vec::new(),
57            tags: Vec::new(),
58        }
59    }
60
61    pub fn to_indented_string(&self, indent: usize) -> String {
62        let mut result = String::new();
63        let prefix = "  ".repeat(indent);
64        result.push_str(&format!(
65            "{}{} {}\n",
66            prefix,
67            "*".repeat(self.level),
68            self.label
69        ));
70
71        for child in &self.children {
72            result.push_str(&child.to_indented_string(indent + 1));
73        }
74
75        result
76    }
77}
78
79impl OrgMode {
80    pub fn new(config: OrgConfig) -> Result<Self, OrgModeError> {
81        let config = config.validate()?;
82
83        Ok(OrgMode { config })
84    }
85
86    pub fn with_defaults() -> Result<Self, OrgModeError> {
87        Self::new(load_org_config(None, None)?)
88    }
89
90    pub fn config(&self) -> &OrgConfig {
91        &self.config
92    }
93}
94
95impl OrgMode {
96    pub fn list_files(
97        &self,
98        tags: Option<&[String]>,
99        limit: Option<usize>,
100    ) -> Result<Vec<String>, OrgModeError> {
101        WalkDir::new(&self.config.org_directory)
102            .into_iter()
103            .filter_map(|entry| match entry {
104                Ok(dir_entry) => {
105                    let path = dir_entry.path();
106
107                    if path.is_file()
108                        && let Some(extension) = path.extension()
109                        && extension == "org"
110                        && let Ok(relative_path) = path.strip_prefix(&self.config.org_directory)
111                        && let Some(path_str) = relative_path.to_str()
112                    {
113                        Some(Ok(path_str.to_string()))
114                    } else {
115                        None
116                    }
117                }
118                Err(e) => Some(Err(OrgModeError::WalkDirError(e))),
119            })
120            .collect::<Result<Vec<String>, OrgModeError>>()
121            .map(|files| {
122                files
123                    .into_iter()
124                    .filter(|path| {
125                        if let Some(tags) = tags {
126                            let file_tags = self.tags_in_file(path).unwrap_or_default();
127                            tags.iter().any(|tag| file_tags.contains(tag))
128                        } else {
129                            true
130                        }
131                    })
132                    .take(limit.unwrap_or(usize::MAX))
133                    .collect::<Vec<String>>()
134            })
135    }
136
137    pub fn search(
138        &self,
139        query: &str,
140        limit: Option<usize>,
141        snippet_max_size: Option<usize>,
142    ) -> Result<Vec<SearchResult>, OrgModeError> {
143        if query.trim().is_empty() {
144            return Ok(vec![]);
145        }
146
147        let mut matcher = Matcher::new(NucleoConfig::DEFAULT);
148        let pattern = Pattern::new(
149            query,
150            CaseMatching::Ignore,
151            Normalization::Smart,
152            AtomKind::Fuzzy,
153        );
154
155        let files = self.list_files(None, None)?;
156        let mut all_results = Vec::new();
157
158        for file in files {
159            let content = match self.read_file(&file) {
160                Ok(content) => content,
161                Err(_) => continue, // Skip files that can't be read
162            };
163
164            let matches = pattern.match_list(
165                content.lines().map(|s| s.to_owned()).collect::<Vec<_>>(),
166                &mut matcher,
167            );
168
169            for (snippet, score) in matches {
170                let snippet = Self::snippet(&snippet, snippet_max_size.unwrap_or(100));
171                all_results.push(SearchResult {
172                    file_path: file.clone(),
173                    snippet,
174                    score,
175                    tags: self.tags_in_file(&file).unwrap_or_default(),
176                });
177            }
178        }
179
180        all_results.sort_by(|a, b| b.score.cmp(&a.score));
181        all_results.truncate(limit.unwrap_or(all_results.len()));
182
183        Ok(all_results)
184    }
185
186    pub fn search_with_tags(
187        &self,
188        query: &str,
189        tags: Option<&[String]>,
190        limit: Option<usize>,
191        snippet_max_size: Option<usize>,
192    ) -> Result<Vec<SearchResult>, OrgModeError> {
193        self.search(query, None, snippet_max_size)
194            .map(|results| match tags {
195                Some(filter_tags) => results
196                    .into_iter()
197                    .filter(|r| filter_tags.iter().any(|tag| r.tags.contains(tag)))
198                    .collect(),
199                None => results,
200            })
201            .map(|mut all_results| {
202                all_results.truncate(limit.unwrap_or(all_results.len()));
203                all_results
204            })
205    }
206
207    pub fn list_files_by_tags(&self, tags: &[String]) -> Result<Vec<String>, OrgModeError> {
208        self.list_files(Some(tags), None)
209    }
210
211    pub fn read_file(&self, path: &str) -> Result<String, OrgModeError> {
212        let org_dir = PathBuf::from(&self.config.org_directory);
213        let full_path = org_dir.join(path);
214
215        if !full_path.exists() {
216            return Err(OrgModeError::IoError(io::Error::new(
217                io::ErrorKind::NotFound,
218                format!("File not found: {}", path),
219            )));
220        }
221
222        if !full_path.is_file() {
223            return Err(OrgModeError::IoError(io::Error::new(
224                io::ErrorKind::InvalidInput,
225                format!("Path is not a file: {}", path),
226            )));
227        }
228
229        fs::read_to_string(full_path).map_err(OrgModeError::IoError)
230    }
231
232    pub fn get_outline(&self, path: &str) -> Result<TreeNode, OrgModeError> {
233        let content = self.read_file(path)?;
234
235        let mut root = TreeNode::new("Document".into());
236        let mut stack: Vec<TreeNode> = Vec::new();
237
238        let mut handler = from_fn(|event| {
239            if let Event::Enter(Container::Headline(h)) = event {
240                let level = h.level();
241                let label = h.title_raw();
242                let tags = h.tags().map(|s| s.to_string()).collect();
243                let node = TreeNode {
244                    label,
245                    level,
246                    tags,
247                    children: Vec::new(),
248                };
249
250                while let Some(n) = stack.last() {
251                    if n.level < level {
252                        break;
253                    }
254                    let finished_node = stack.pop().unwrap();
255                    if let Some(parent) = stack.last_mut() {
256                        parent.children.push(finished_node);
257                    } else {
258                        root.children.push(finished_node);
259                    }
260                }
261
262                stack.push(node);
263            }
264        });
265
266        Org::parse(&content).traverse(&mut handler);
267
268        while let Some(node) = stack.pop() {
269            if let Some(parent) = stack.last_mut() {
270                parent.children.push(node);
271            } else {
272                root.children.push(node);
273            }
274        }
275
276        Ok(root)
277    }
278
279    pub fn get_heading(&self, path: &str, heading: &str) -> Result<String, OrgModeError> {
280        let content = self.read_file(path)?;
281
282        let heading_path: Vec<&str> = heading.split('/').collect();
283        let mut current_level = 0;
284        let mut found = None;
285
286        let mut handler = from_fn_with_ctx(|event, ctx| {
287            if let Event::Enter(Container::Headline(h)) = event {
288                let title = h.title_raw();
289                let level = h.level();
290
291                if let Some(part) = heading_path.get(current_level) {
292                    if title == *part {
293                        if level == heading_path.len() {
294                            found = Some(h);
295                            ctx.stop();
296                        }
297                        current_level += 1;
298                    }
299                } else {
300                    ctx.stop()
301                }
302            }
303        });
304
305        Org::parse(&content).traverse(&mut handler);
306
307        found
308            .ok_or_else(|| OrgModeError::InvalidHeadingPath(heading.into()))
309            .map(|h| h.raw())
310    }
311
312    pub fn get_element_by_id(&self, id: &str) -> Result<String, OrgModeError> {
313        let files = self.list_files(None, None)?;
314
315        let found = files.iter().find_map(|path| {
316            self.read_file(path)
317                .map(|content| self.search_id(content, id))
318                .unwrap_or_default()
319        });
320
321        found.ok_or_else(|| OrgModeError::InvalidElementId(id.into()))
322    }
323
324    fn search_id(&self, content: String, id: &str) -> Option<String> {
325        let mut found = None;
326        let has_id_property = |properties: Option<PropertyDrawer>| {
327            properties
328                .and_then(|props| {
329                    props
330                        .to_hash_map()
331                        .into_iter()
332                        .find(|(k, v)| k.to_lowercase() == "id" && v == id)
333                })
334                .is_some()
335        };
336        let mut handler = from_fn_with_ctx(|event, ctx| {
337            if let Event::Enter(Container::Headline(ref h)) = event
338                && has_id_property(h.properties())
339            {
340                found = Some(h.raw());
341                ctx.stop();
342            } else if let Event::Enter(Container::Document(ref d)) = event
343                && has_id_property(d.properties())
344            {
345                found = Some(d.raw());
346                ctx.stop();
347            }
348        });
349
350        Org::parse(&content).traverse(&mut handler);
351
352        found
353    }
354
355    fn snippet(s: &str, max: usize) -> String {
356        if s.chars().count() > max {
357            s.chars().take(max).collect::<String>() + "..."
358        } else {
359            s.to_string()
360        }
361    }
362
363    fn tags_in_file(&self, path: &str) -> Result<Vec<String>, OrgModeError> {
364        let content = self.read_file(path)?;
365        let mut tags = Vec::new();
366
367        let mut handler = from_fn(|event| {
368            if let Event::Enter(Container::Headline(h)) = event {
369                tags.extend(h.tags().map(|s| s.to_string()));
370            }
371        });
372
373        Org::parse(&content).traverse(&mut handler);
374
375        Ok(tags)
376    }
377}