2 files changed, 89 insertions, 27 deletions

diff --git a/src/constants.rs b/src/constants.rs
index 2cdec2d..6536f63 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -3,3 +3,6 @@ pub const REPLACEMENT_CHAPTER: &str = "lpm::new_chapter_name";
 pub const REPLACEMENT_CHAPTER_SECTION: &str = "lpm::current_chapter_name::title_case";
 pub const REPLACMENT_SECTION_TITLE: &str = "lpm::new_section_name";
 pub const DATE: &str = "lpm::current_date";
+
+pub const NEXT_CHAPTER: &str = "% lpm::next_chapter_marker";
+pub const NEXT_CHAPTER_INCLUDE_ONLY: &str = "% lpm::next_chapter_includeonly_marker";
diff --git a/src/new/chapter.rs b/src/new/chapter.rs
index 7628db6..9fb5301 100644
--- a/src/new/chapter.rs
+++ b/src/new/chapter.rs
@@ -1,9 +1,11 @@
 use std::{fmt::Display, fs, path::Path};
 
 use anyhow::{Context, Result};
+use log::error;
 
 use crate::{
     config_file::Config,
+    constants::{NEXT_CHAPTER, NEXT_CHAPTER_INCLUDE_ONLY},
     file_tree::{FileTree, GeneratedFile},
 };
 
@@ -141,6 +143,30 @@ fn new_chapter_file(
     )
 }
 
+fn find(input: &str, marker: &str, source_path: &Path) -> anyhow::Result<usize> {
+    let find_index = input.find(marker).with_context(|| {
+        format!(
+            "Failed to find the '{}' in '{}'.",
+            marker,
+            source_path.display()
+        )
+    })?;
+    Ok(find_index)
+}
+
+fn find_and_append(
+    input: &mut String,
+    marker: &str,
+    source_path: &Path,
+    insertion: &str,
+) -> anyhow::Result<()> {
+    let find_index = find(input, marker, source_path)?;
+
+    input.insert_str(find_index, insertion);
+
+    Ok(())
+}
+
 fn new_main_file(
     project_root: &Path,
     config: &Config,
@@ -151,38 +177,71 @@ fn new_main_file(
     let main_path = project_root.join(&config.main_file);
     let mut main_text = fs::read_to_string(&main_path)?;
 
-    let chapter_includeonly: String = format!(
-        "\\includeonly{{content/{}/{}}}",
+    let chapter_path = format!(
+        "content/{}/{}",
         ChapterName::from_str(name, last_chapter_number).to_string(),
         "chapter.tex",
     );
 
-    if &last_chapter_name.as_str() == &"static" && last_chapter_number == 0 {
-        // This is the first added chapter; The `\includeonly` will be empty.
-        main_text = main_text.replace("\\includeonly{}", &chapter_includeonly)
-    } else {
-        main_text = main_text.replace(
-            &format!(
-                "\\includeonly{{content/{}/{}}}",
-                ChapterName::new(last_chapter_name, last_chapter_number - 1).to_string(),
-                "chapter.tex",
-            ),
-            &chapter_includeonly,
-        )
-    };
-
-    let find_index = main_text
-        .find("% NEXT_CHAPTER")
-        .expect("The % NEXT_CHAPTER maker must exist");
-
-    main_text.insert_str(
-        find_index,
-        &format!(
-            "\\include{{content/{}/{}}}\n    ",
-            ChapterName::from_str(name, last_chapter_number).to_string(),
+    // Check if this is the first added chapter; Then there should be no other path before this one.
+    if last_chapter_name.as_str() != "static" && last_chapter_number != 0 {
+        let old_path = format!(
+            "content/{}/{}",
+            ChapterName::new(last_chapter_name, last_chapter_number - 1).to_string(),
             "chapter.tex",
-        ),
-    );
+        );
+        let find_index = find(&main_text, NEXT_CHAPTER_INCLUDE_ONLY, &main_path)?;
+
+        // The last element adds the indentation.
+        let previous_includeonly_white_space = main_text
+            .as_bytes()
+            .iter()
+            .take(find_index)
+            .rev()
+            .take_while(|byte| **byte == ' ' as u8 || **byte == '\n' as u8)
+            .count();
+        let previous_includeonly = find_index - (old_path.len() + previous_includeonly_white_space);
+
+        let old_include = {
+            let old_bytes = main_text
+                .as_bytes()
+                .into_iter()
+                .take(find_index)
+                .rev()
+                .take(find_index - previous_includeonly)
+                .rev()
+                .map(|b| *b)
+                .collect::<Vec<u8>>();
+
+            let string = String::from_utf8(old_bytes).expect("This should be valid utf8");
+
+            string.trim().to_owned()
+        };
+
+        if old_include != old_path.as_str() {
+            error!(
+                "Failed to determine the old includeonly path: Found '{}' but expected '{}'. \
+                 Refusing to add a comment to it.",
+                old_include, old_path
+            )
+        } else {
+            main_text.insert_str(previous_includeonly, "% ");
+        }
+    }
+
+    find_and_append(
+        &mut main_text,
+        NEXT_CHAPTER_INCLUDE_ONLY,
+        &main_path,
+        &format!("{}\n    ", chapter_path),
+    )?;
+
+    find_and_append(
+        &mut main_text,
+        NEXT_CHAPTER,
+        &main_path,
+        &format!("\\include{{{}}}\n", chapter_path),
+    )?;
 
     Ok(GeneratedFile::new(main_path, main_text))
 }