1use indoc::indoc;
2use log::debug;
3use std::{collections::HashMap, env, fs::read_to_string};
4use toml_edit::{value, DocumentMut, Item};
5
6use serde::{Deserialize, Serialize};
7
8use crate::graph::GraphContext;
9
10use super::{node::NodeIter, NodeId};
11
12const CONFIG_FILE_NAME: &str = "config.toml";
13const IWE_MARKER: &str = ".iwe";
14
15pub const DEFAULT_KEY_DATE_FORMAT: &str = "%Y-%m-%d";
16
17#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
18pub struct MarkdownOptions {
19 pub refs_extension: String,
20 pub date_format: Option<String>,
21}
22
23impl Default for MarkdownOptions {
24 fn default() -> Self {
25 Self {
26 refs_extension: String::new(),
27 date_format: Some("%b %d, %Y".into()),
28 }
29 }
30}
31
32#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
33pub struct LibraryOptions {
34 pub path: String,
35 pub date_format: Option<String>,
36 pub prompt_key_prefix: Option<String>,
37 pub default_template: Option<String>,
38}
39
40#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)]
41pub struct CompletionOptions {
42 pub link_format: Option<LinkType>,
43}
44
45impl Default for LibraryOptions {
46 fn default() -> Self {
47 Self {
48 path: String::new(),
49 date_format: Some(DEFAULT_KEY_DATE_FORMAT.into()),
50 prompt_key_prefix: None,
51 default_template: None,
52 }
53 }
54}
55
56#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
57pub struct Configuration {
58 pub version: Option<u32>,
59 pub markdown: MarkdownOptions,
60 pub library: LibraryOptions,
61 #[serde(default)]
62 pub completion: CompletionOptions,
63 pub models: HashMap<String, Model>,
64 pub actions: HashMap<String, ActionDefinition>,
65 #[serde(default)]
66 pub templates: HashMap<String, NoteTemplate>,
67}
68
69#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)]
70pub struct Model {
71 pub api_key_env: String,
72 pub base_url: String,
73
74 pub name: String,
75 pub max_tokens: Option<u64>,
76 pub max_completion_tokens: Option<u64>,
77 pub temperature: Option<f32>,
78}
79
80#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
81#[serde(tag = "type")]
82pub enum ActionDefinition {
83 #[serde(rename = "transform")]
84 Transform(Transform),
85 #[serde(rename = "attach")]
86 Attach(Attach),
87 #[serde(rename = "sort")]
88 Sort(Sort),
89 #[serde(rename = "inline")]
90 Inline(Inline),
91 #[serde(rename = "extract")]
92 Extract(Extract),
93 #[serde(rename = "extract_all")]
94 ExtractAll(ExtractAll),
95 #[serde(rename = "link")]
96 Link(Link),
97}
98
99#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
100pub struct Transform {
101 pub title: String,
102 pub model: String,
103 pub prompt_template: String,
104 pub context: Context,
105}
106
107#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
108pub struct Attach {
109 pub title: String,
110 pub key_template: String,
111 pub document_template: String,
112}
113
114#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
115pub struct Sort {
116 pub title: String,
117 pub reverse: Option<bool>,
118}
119
120#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
121pub struct Inline {
122 pub title: String,
123 pub inline_type: InlineType,
124 pub keep_target: Option<bool>,
125}
126
127#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
128pub enum InlineType {
129 #[serde(rename = "section")]
130 Section,
131 #[serde(rename = "quote")]
132 Quote,
133}
134
135#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
136pub struct Extract {
137 pub title: String,
138 pub link_type: Option<LinkType>,
139 pub key_template: String,
140}
141
142#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
143pub struct ExtractAll {
144 pub title: String,
145 pub link_type: Option<LinkType>,
146 pub key_template: String,
147}
148
149#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
150pub struct Link {
151 pub title: String,
152 pub link_type: Option<LinkType>,
153 pub key_template: String,
154}
155
156#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
157pub enum LinkType {
158 #[serde(rename = "markdown")]
159 Markdown,
160 #[serde(rename = "wiki")]
161 WikiLink,
162}
163
164#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
165pub struct NoteTemplate {
166 pub key_template: String,
167 pub document_template: String,
168}
169
170impl Default for Configuration {
171 fn default() -> Self {
172 Self {
173 version: Some(1),
174 markdown: Default::default(),
175 library: Default::default(),
176 completion: Default::default(),
177 models: Default::default(),
178 actions: Default::default(),
179 templates: Default::default(),
180 }
181 }
182}
183
184impl Configuration {
185 pub fn template() -> Self {
186 let mut template = Self {
187 version: Some(2),
188 ..Default::default()
189 };
190
191 template.models.insert(
192 "default".into(),
193 Model {
194 api_key_env: "OPENAI_API_KEY".to_string(),
195 base_url: "https://api.openai.com".to_string(),
196 name: "gpt-4o".into(),
197 max_tokens: None,
198 max_completion_tokens: None,
199 temperature: None,
200 },
201 );
202
203 template.models.insert(
204 "fast".into(),
205 Model {
206 api_key_env: "OPENAI_API_KEY".into(),
207 base_url: "https://api.openai.com".into(),
208 name: "gpt-4o-mini".to_string(),
209 max_tokens: None,
210 max_completion_tokens: None,
211 temperature: None,
212 },
213 );
214
215 template.actions.insert(
216 "today".into(),
217 ActionDefinition::Attach(Attach {
218 title: "Add Date".into(),
219 key_template: "{{today}}".into(),
220 document_template: "# {{today}}\n\n{{content}}\n".into(),
221 }),
222 );
223
224 template.actions.insert(
225 "rewrite".into(),
226 ActionDefinition::Transform(
227 Transform {
228 title: "Rewrite".into(),
229 model: "default".into(),
230 prompt_template: indoc! {r##"
231 Here's a text that I'm going to ask you to edit. The text is marked with {{context_start}}{{context_end}} tag.
232
233 The part you'll need to update is marked with {{update_start}}{{update_end}}.
234
235 {{context_start}}
236
237 {{context}}
238
239 {{context_end}}
240
241 - You can't replace entire text, your answer will be inserted in place of the {{update_start}}{{update_end}}. Don't include the {{context_start}}{{context_end}} and {{context_start}}{{context_end}} tags in your output.
242 - Preserve the links in the text. Do not return list item "-" or header "#" prefix
243
244 Your goal is to rewrite a given text to improve its clarity and readability. Ensure the language remains personable and not overly formal. Focus on simplifying language, organizing sentences logically, and removing ambiguity while maintaining a conversational tone.
245 "##}.to_string(),
246 context: Context::Document
247 }
248 ),
249 );
250
251 template.actions.insert (
252 "expand".to_string(),
253 ActionDefinition::Transform(
254 Transform {
255 title: "Expand".to_string(),
256 model: "default".to_string(),
257 prompt_template: indoc! {r##"
258 Here's a text that I'm going to ask you to edit. The text is marked with {{context_start}}{{context_end}} tag.
259
260 The part you'll need to update is marked with {{update_start}}{{update_end}}.
261
262 {{context_start}}
263
264 {{context}}
265
266 {{context_end}}
267
268 - You can't replace entire text, your answer will be inserted in place of the {{update_start}}{{update_end}}. Don't include the {{context_start}}{{context_end}} and {{context_start}}{{context_end}} tags in your output.
269 - Preserve the links in the text. Do not return list item "-" or header "#" prefix
270
271 Expand the text you need to update, generate a couple paragraphs.
272 "##}.to_string(),
273 context: Context::Document
274 }
275 ),
276 );
277
278 template.actions.insert (
279 "keywords".into(),
280 ActionDefinition::Transform(
281 Transform {
282 title: "Keywords".to_string(),
283 model: "default".to_string(),
284 prompt_template: indoc! {r##"
285 Here's a text that I'm going to ask you to edit. The text is marked with {{context_start}}{{context_end}} tag.
286
287 The part you'll need to update is marked with {{update_start}}{{update_end}}.
288
289 {{context_start}}
290
291 {{context}}
292
293 {{context_end}}
294
295 - You can't replace entire text, your answer will be inserted in place of the {{update_start}}{{update_end}}. Don't include the {{context_start}}{{context_end}} and {{context_start}}{{context_end}} tags in your output.
296
297 Mark most important keywords with bold using ** markdown syntax. Keep the text unchanged!
298 "##}.to_string(),
299 context: Context::Document
300 }
301 ),
302 );
303
304 template.actions.insert(
305 "emoji".into(),
306 ActionDefinition::Transform(
307 Transform {
308 title: "Emojify".to_string(),
309 model: "default".to_string(),
310 prompt_template: indoc! {r##"
311 Here's a text that I'm going to ask you to edit. The text is marked with {{context_start}} {{context_end}} tags.
312
313 - The part you'll need to update is marked with {{update_start}} {{update_end}} tags.
314 - You can't replace entire text, your answer will be inserted in between {{update_start}} {{update_end}} tags.
315 - Add a relevant emoji one per list item (prior to list item text), header (prior to header text) or paragraph. Keep the text otherwise unchanged.
316 - Don't include the {{update_start}} {{update_end}} tags in your answer.
317
318 {{context_start}}
319
320 {{context}}
321
322 {{context_end}}
323 "##}.to_string(),
324 context: Context::Document
325 }
326 )
327 );
328
329 template.actions.insert(
330 "sort".into(),
331 ActionDefinition::Sort(Sort {
332 title: "Sort A-Z".into(),
333 reverse: Some(false),
334 }),
335 );
336
337 template.actions.insert(
338 "sort_desc".into(),
339 ActionDefinition::Sort(Sort {
340 title: "Sort Z-A".into(),
341 reverse: Some(true),
342 }),
343 );
344
345 template.actions.insert(
346 "inline_section".into(),
347 ActionDefinition::Inline(Inline {
348 title: "Inline section".into(),
349 inline_type: InlineType::Section,
350 keep_target: Some(false),
351 }),
352 );
353
354 template.actions.insert(
355 "inline_quote".into(),
356 ActionDefinition::Inline(Inline {
357 title: "Inline quote".into(),
358 inline_type: InlineType::Quote,
359 keep_target: Some(false),
360 }),
361 );
362
363 template.actions.insert(
364 "extract".into(),
365 ActionDefinition::Extract(Extract {
366 title: "Extract".into(),
367 link_type: Some(LinkType::Markdown),
368 key_template: "{{id}}".into(),
369 }),
370 );
371
372 template.actions.insert(
373 "extract_all".into(),
374 ActionDefinition::ExtractAll(ExtractAll {
375 title: "Extract all subsections".into(),
376 link_type: Some(LinkType::Markdown),
377 key_template: "{{id}}".into(),
378 }),
379 );
380
381 template.actions.insert(
382 "link".into(),
383 ActionDefinition::Link(Link {
384 title: "Link".into(),
385 link_type: Some(LinkType::Markdown),
386 key_template: "{{id}}".into(),
387 }),
388 );
389
390 template.templates.insert(
391 "default".into(),
392 NoteTemplate {
393 key_template: "{{slug}}".into(),
394 document_template: "# {{title}}\n\n{{content}}".into(),
395 },
396 );
397
398 template
399 }
400}
401
402#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
403pub enum Context {
404 Document,
405}
406
407#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
408pub enum Operation {
409 Replace,
410}
411
412#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
413pub enum TargetType {
414 Block,
415 Paragraph,
416 List,
417}
418
419impl TargetType {
420 pub fn acceptable_target(&self, id: NodeId, context: impl GraphContext) -> Option<NodeId> {
421 match self {
422 TargetType::Block => Some(id).filter(|id| !context.node(*id).is_section()),
423 TargetType::Paragraph => Some(id).filter(|id| context.node(*id).is_leaf()),
424 TargetType::List => Some(id).filter(|id| context.node(*id).is_leaf()),
425 }
426 }
427}
428
429pub fn load_config() -> Configuration {
430 let current_dir = env::current_dir().expect("to get current dir");
431 let mut config_path = current_dir.clone();
432 config_path.push(IWE_MARKER);
433 config_path.push(CONFIG_FILE_NAME);
434
435 if config_path.exists() {
436 debug!("reading config from path: {:?}", config_path);
437
438 let configuration = migrate(&read_to_string(config_path).expect("to read config file"));
439
440 toml::from_str::<Configuration>(&configuration).expect("to parse config file")
441 } else {
442 debug!("using default configuration");
443 Configuration::template()
444 }
445}
446
447fn migrate(config: &str) -> String {
448 let doc = config.parse::<DocumentMut>().expect("valid TOML");
449 let current_version = doc
450 .get("version")
451 .and_then(|v| v.as_value())
452 .and_then(|v| v.as_integer())
453 .unwrap_or(0);
454
455 let mut updated = config.to_string();
456 let mut needs_update = false;
457
458 if current_version < 1 {
460 debug!("applying migrations from version 0 to 1");
461 updated = add_default_type_to_actions(&updated);
462 updated = add_default_code_actions(&updated);
463 updated = set_config_version(&updated, 1);
464 needs_update = true;
465 }
466
467 if current_version < 2 {
469 debug!("applying migrations from version 1 to 2");
470 updated = add_refs_extension_field(&updated);
471 updated = add_link_action(&updated);
472 updated = set_config_version(&updated, 2);
473 needs_update = true;
474 }
475
476 if needs_update {
477 debug!("configuration file migration applied");
478 let current_dir = env::current_dir().expect("to get current dir");
479 let mut config_path = current_dir.clone();
480 config_path.push(IWE_MARKER);
481 config_path.push(CONFIG_FILE_NAME);
482
483 debug!("updating configuration file");
484 std::fs::write(config_path, &updated).expect("to write updated config file");
485 }
486
487 updated
488}
489
490fn add_default_type_to_actions(input: &str) -> String {
491 let mut doc = input.parse::<DocumentMut>().expect("valid TOML");
492
493 if let Some(Item::Table(actions)) = doc.get_mut("actions") {
494 for (_, action) in actions.iter_mut() {
495 if let Item::Table(action_table) = action {
496 action_table.entry("type").or_insert(value("transform"));
497 }
498 }
499 }
500
501 doc.to_string()
502}
503
504fn add_refs_extension_field(input: &str) -> String {
505 let mut doc = input.parse::<DocumentMut>().expect("valid TOML");
506
507 if doc.get("markdown").is_none() {
508 doc["markdown"] = Item::Table(toml_edit::Table::new());
509 }
510
511 if let Some(Item::Table(markdown)) = doc.get_mut("markdown") {
512 if markdown.get("refs_extension").is_none() {
513 markdown.insert("refs_extension", value(""));
514 }
515 }
516
517 doc.to_string()
518}
519
520fn add_link_action(input: &str) -> String {
521 let mut doc = input.parse::<DocumentMut>().expect("valid TOML");
522
523 if doc.get("actions").is_none() {
524 doc["actions"] = Item::Table(toml_edit::Table::new());
525 }
526
527 if let Some(Item::Table(actions)) = doc.get_mut("actions") {
528 let has_link = actions.iter().any(|(_, action)| {
530 if let Item::Table(action_table) = action {
531 if let Some(Item::Value(action_type)) = action_table.get("type") {
532 if let Some(type_str) = action_type.as_str() {
533 return type_str == "link";
534 }
535 }
536 }
537 false
538 });
539
540 if !has_link {
541 let mut link_table = toml_edit::Table::new();
542 link_table.insert("type", value("link"));
543 link_table.insert("title", value("Link word"));
544 link_table.insert("link_type", value("markdown"));
545 link_table.insert("key_template", value("{{id}}"));
546 actions.insert("link", Item::Table(link_table));
547 }
548 }
549
550 doc.to_string()
551}
552
553fn set_config_version(input: &str, version: i64) -> String {
554 let mut doc = input.parse::<DocumentMut>().expect("valid TOML");
555
556 doc.insert("version", value(version));
557
558 doc.to_string()
559}
560
561fn add_default_code_actions(input: &str) -> String {
562 let mut doc = input.parse::<DocumentMut>().expect("valid TOML");
563
564 if doc.get("actions").is_none() {
565 doc["actions"] = Item::Table(toml_edit::Table::new());
566 }
567
568 if let Some(Item::Table(actions)) = doc.get_mut("actions") {
569 let mut has_extract = false;
570 let mut has_extract_all = false;
571 let mut has_inline = false;
572
573 for (_, action) in actions.iter() {
574 if let Item::Table(action_table) = action {
575 if let Some(Item::Value(action_type)) = action_table.get("type") {
576 if let Some(type_str) = action_type.as_str() {
577 match type_str {
578 "extract" => has_extract = true,
579 "extract_all" => has_extract_all = true,
580 "inline" => has_inline = true,
581 _ => {}
582 }
583 }
584 }
585 }
586 }
587
588 if !has_extract {
589 let mut extract_table = toml_edit::Table::new();
590 extract_table.insert("type", value("extract"));
591 extract_table.insert("title", value("Extract"));
592 extract_table.insert("link_type", value("markdown"));
593 extract_table.insert("key_template", value("{{id}}"));
594 actions.insert("extract", Item::Table(extract_table));
595 }
596
597 if !has_extract_all {
598 let mut extract_all_table = toml_edit::Table::new();
599 extract_all_table.insert("type", value("extract_all"));
600 extract_all_table.insert("title", value("Extract all subsections"));
601 extract_all_table.insert("link_type", value("markdown"));
602 extract_all_table.insert("key_template", value("{{id}}"));
603 actions.insert("extract_all", Item::Table(extract_all_table));
604 }
605
606 if !has_inline {
607 let mut inline_section_table = toml_edit::Table::new();
608 inline_section_table.insert("type", value("inline"));
609 inline_section_table.insert("title", value("Inline section"));
610 inline_section_table.insert("inline_type", value("section"));
611 inline_section_table.insert("keep_target", value(false));
612 actions.insert("inline_section", Item::Table(inline_section_table));
613
614 let mut inline_quote_table = toml_edit::Table::new();
615 inline_quote_table.insert("type", value("inline"));
616 inline_quote_table.insert("title", value("Inline quote"));
617 inline_quote_table.insert("inline_type", value("quote"));
618 inline_quote_table.insert("keep_target", value(false));
619 actions.insert("inline_quote", Item::Table(inline_quote_table));
620 }
621 }
622
623 doc.to_string()
624}