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, };
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}