weave_content/

writeback.rs

use std::path::Path;

/// A generated ID that needs to be written back to a source file.
#[derive(Debug)]
pub struct PendingId {
    /// 1-indexed line number where the ID should be inserted.
    pub line: usize,
    /// The generated NULID string.
    pub id: String,
    /// What kind of insertion to perform.
    pub kind: WriteBackKind,
}

/// The type of write-back insertion.
#[derive(Debug)]
pub enum WriteBackKind {
    /// Insert `id: <NULID>` into YAML front matter of an entity file.
    /// `line` is the line of the `---` closing delimiter; insert before it.
    EntityFrontMatter,
    /// Insert `id: <NULID>` into YAML front matter of a case file.
    /// `line` is the line of the `---` closing delimiter; insert before it.
    CaseId,
    /// Insert `- id: <NULID>` bullet after the H3 heading line.
    /// `line` is the line of the `### Name` heading.
    InlineEvent,
    /// Insert `  - id: <NULID>` nested bullet after the relationship line.
    /// `line` is the line of `- Source -> Target: type`.
    Relationship,
    /// Insert `  id: <NULID>` after a `## Related Cases` bullet.
    /// `line` is the line of `- case/path`.
    RelatedCase,
    /// Insert `  id: <NULID>` after a `## Involved` bullet.
    /// `line` is the line of `- Entity Name` (0 if section doesn't exist yet).
    /// `entity_name` is needed to create the bullet when the section is missing.
    InvolvedIn { entity_name: String },
    /// Insert `  id: <NULID>` after a `## Timeline` bullet.
    /// `line` is the line of `- Event A -> Event B`.
    TimelineEdge,
}

/// Apply pending ID write-backs to a file's content.
///
/// Insertions are applied from bottom to top (highest line number first)
/// so that earlier insertions don't shift line numbers for later ones.
///
/// Returns the modified content, or None if no changes were needed.
pub fn apply_writebacks(content: &str, pending: &mut [PendingId]) -> Option<String> {
    if pending.is_empty() {
        return None;
    }

    // Separate out InvolvedIn entries with line==0 (need section creation)
    // vs normal insertions
    let mut new_involved: Vec<(String, String)> = Vec::new(); // (entity_name, id)
    let mut normal: Vec<&PendingId> = Vec::new();

    for p in pending.iter() {
        if let WriteBackKind::InvolvedIn { ref entity_name } = p.kind
            && p.line == 0
        {
            new_involved.push((entity_name.clone(), p.id.clone()));
            continue;
        }
        normal.push(p);
    }

    // Sort normal insertions by line descending so we insert from bottom to top
    normal.sort_by_key(|p| std::cmp::Reverse(p.line));

    let mut lines: Vec<String> = content.lines().map(String::from).collect();
    // Preserve trailing newline if original had one
    let trailing_newline = content.ends_with('\n');

    for p in &normal {
        let insert_text = match p.kind {
            WriteBackKind::EntityFrontMatter | WriteBackKind::CaseId => {
                format!("id: {}", p.id)
            }
            WriteBackKind::InlineEvent => {
                format!("- id: {}", p.id)
            }
            WriteBackKind::Relationship => {
                format!("  - id: {}", p.id)
            }
            WriteBackKind::RelatedCase
            | WriteBackKind::InvolvedIn { .. }
            | WriteBackKind::TimelineEdge => {
                format!("  id: {}", p.id)
            }
        };

        let insert_after = match p.kind {
            WriteBackKind::EntityFrontMatter | WriteBackKind::CaseId => {
                // Check if there's an existing empty `id:` line in front matter
                // and replace it instead of inserting a duplicate.
                let fm_end_idx = p.line.saturating_sub(1); // 0-indexed
                let mut replaced = false;
                for i in 0..fm_end_idx.min(lines.len()) {
                    let trimmed = lines[i].trim();
                    if trimmed == "id:" || trimmed == "id: " {
                        lines[i] = insert_text.clone();
                        replaced = true;
                        break;
                    }
                }
                if !replaced && fm_end_idx <= lines.len() {
                    lines.insert(fm_end_idx, insert_text);
                }
                continue;
            }
            WriteBackKind::InlineEvent
            | WriteBackKind::Relationship
            | WriteBackKind::RelatedCase
            | WriteBackKind::InvolvedIn { .. }
            | WriteBackKind::TimelineEdge => {
                // Insert after the heading/relationship/bullet line.
                p.line // 1-indexed line; inserting after means index = line (0-indexed + 1)
            }
        };

        if insert_after <= lines.len() {
            lines.insert(insert_after, insert_text);
        }
    }

    // Append ## Involved section for entities without existing entries
    if !new_involved.is_empty() {
        // Find insertion point: before ## Timeline or ## Related Cases, or at end
        let insert_idx = find_section_insert_point(&lines);

        let mut section_lines = Vec::new();
        // Add blank line before section if needed
        if insert_idx > 0 && !lines[insert_idx - 1].is_empty() {
            section_lines.push(String::new());
        }
        section_lines.push("## Involved".to_string());
        section_lines.push(String::new());
        for (name, id) in &new_involved {
            section_lines.push(format!("- {name}"));
            section_lines.push(format!("  id: {id}"));
        }

        // Insert all section lines at the insertion point
        for (offset, line) in section_lines.into_iter().enumerate() {
            lines.insert(insert_idx + offset, line);
        }
    }

    let mut result = lines.join("\n");
    if trailing_newline {
        result.push('\n');
    }
    Some(result)
}

/// Find the best insertion point for a new `## Involved` section.
/// Prefers inserting before `## Timeline` or `## Related Cases`.
/// Falls back to end of file.
fn find_section_insert_point(lines: &[String]) -> usize {
    for (i, line) in lines.iter().enumerate() {
        let trimmed = line.trim();
        if trimmed == "## Timeline" || trimmed == "## Related Cases" {
            return i;
        }
    }
    lines.len()
}

/// Write modified content back to the file at `path`.
pub fn write_file(path: &Path, content: &str) -> Result<(), String> {
    std::fs::write(path, content)
        .map_err(|e| format!("{}: error writing file: {e}", path.display()))
}

/// Find the line number of the closing `---` in YAML front matter.
/// Returns None if the file has no front matter or the closing delimiter
/// is not found.
pub fn find_front_matter_end(content: &str) -> Option<usize> {
    let mut in_front_matter = false;
    for (i, line) in content.lines().enumerate() {
        let trimmed = line.trim();
        if trimmed == "---" && !in_front_matter {
            in_front_matter = true;
        } else if trimmed == "---" && in_front_matter {
            return Some(i + 1); // 1-indexed
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn writeback_entity_front_matter() {
        let content = "---\n---\n\n# Mark Bonnick\n\n- nationality: British\n";
        let end_line = find_front_matter_end(content).unwrap();
        assert_eq!(end_line, 2);

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JXYZ123456789ABCDEFGHIJK".to_string(),
            kind: WriteBackKind::EntityFrontMatter,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        assert!(result.contains("id: 01JXYZ123456789ABCDEFGHIJK"));
        // id should appear between the two ---
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[0], "---");
        assert_eq!(lines[1], "id: 01JXYZ123456789ABCDEFGHIJK");
        assert_eq!(lines[2], "---");
    }

    #[test]
    fn writeback_entity_front_matter_with_existing_fields() {
        let content = "---\nother: value\n---\n\n# Test\n";
        let end_line = find_front_matter_end(content).unwrap();
        assert_eq!(end_line, 3);

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JABC000000000000000000AA".to_string(),
            kind: WriteBackKind::EntityFrontMatter,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[0], "---");
        assert_eq!(lines[1], "other: value");
        assert_eq!(lines[2], "id: 01JABC000000000000000000AA");
        assert_eq!(lines[3], "---");
    }

    #[test]
    fn writeback_inline_event() {
        let content = "\
## Events

### Dismissal
- occurred_at: 2024-12-24
- event_type: termination
";
        // H3 heading is on line 3 (1-indexed)
        let mut pending = vec![PendingId {
            line: 3,
            id: "01JXYZ123456789ABCDEFGHIJK".to_string(),
            kind: WriteBackKind::InlineEvent,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[2], "### Dismissal");
        assert_eq!(lines[3], "- id: 01JXYZ123456789ABCDEFGHIJK");
        assert_eq!(lines[4], "- occurred_at: 2024-12-24");
    }

    #[test]
    fn writeback_relationship() {
        let content = "\
## Relationships

- Alice -> Bob: employed_by
  - source: https://example.com
";
        // Relationship line is on line 3 (1-indexed)
        let mut pending = vec![PendingId {
            line: 3,
            id: "01JXYZ123456789ABCDEFGHIJK".to_string(),
            kind: WriteBackKind::Relationship,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[2], "- Alice -> Bob: employed_by");
        assert_eq!(lines[3], "  - id: 01JXYZ123456789ABCDEFGHIJK");
        assert_eq!(lines[4], "  - source: https://example.com");
    }

    #[test]
    fn writeback_multiple_insertions() {
        let content = "\
## Events

### Event A
- occurred_at: 2024-01-01

### Event B
- occurred_at: 2024-06-01

## Relationships

- Event A -> Event B: associate_of
";
        let mut pending = vec![
            PendingId {
                line: 3,
                id: "01JAAA000000000000000000AA".to_string(),
                kind: WriteBackKind::InlineEvent,
            },
            PendingId {
                line: 6,
                id: "01JBBB000000000000000000BB".to_string(),
                kind: WriteBackKind::InlineEvent,
            },
            PendingId {
                line: 10,
                id: "01JCCC000000000000000000CC".to_string(),
                kind: WriteBackKind::Relationship,
            },
        ];

        let result = apply_writebacks(content, &mut pending).unwrap();
        assert!(result.contains("- id: 01JAAA000000000000000000AA"));
        assert!(result.contains("- id: 01JBBB000000000000000000BB"));
        assert!(result.contains("  - id: 01JCCC000000000000000000CC"));
    }

    #[test]
    fn writeback_empty_pending() {
        let content = "some content\n";
        let mut pending = Vec::new();
        assert!(apply_writebacks(content, &mut pending).is_none());
    }

    #[test]
    fn writeback_preserves_trailing_newline() {
        let content = "---\n---\n\n# Test\n";
        let mut pending = vec![PendingId {
            line: 2,
            id: "01JABC000000000000000000AA".to_string(),
            kind: WriteBackKind::EntityFrontMatter,
        }];
        let result = apply_writebacks(content, &mut pending).unwrap();
        assert!(result.ends_with('\n'));
    }

    #[test]
    fn writeback_case_id() {
        let content = "---\nsources:\n  - https://example.com\n---\n\n# Some Case\n";
        let end_line = find_front_matter_end(content).unwrap();
        assert_eq!(end_line, 4);

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JXYZ123456789ABCDEFGHIJK".to_string(),
            kind: WriteBackKind::CaseId,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[0], "---");
        assert_eq!(lines[3], "id: 01JXYZ123456789ABCDEFGHIJK");
        assert_eq!(lines[4], "---");
        assert!(!result.contains("\nnulid:"));
    }

    #[test]
    fn find_front_matter_end_basic() {
        assert_eq!(find_front_matter_end("---\nid: foo\n---\n"), Some(3));
        assert_eq!(find_front_matter_end("---\n---\n"), Some(2));
        assert_eq!(find_front_matter_end("no front matter"), None);
        assert_eq!(find_front_matter_end("---\nunclosed"), None);
    }

    #[test]
    fn writeback_related_case() {
        let content = "\
## Related Cases

- id/corruption/2013/some-case
  description: Related scandal
";
        // Bullet line is on line 3 (1-indexed)
        let mut pending = vec![PendingId {
            line: 3,
            id: "01JREL000000000000000000AA".to_string(),
            kind: WriteBackKind::RelatedCase,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[2], "- id/corruption/2013/some-case");
        assert_eq!(lines[3], "  id: 01JREL000000000000000000AA");
        assert_eq!(lines[4], "  description: Related scandal");
    }

    #[test]
    fn writeback_involved_in_existing_section() {
        let content = "\
## Involved

- John Doe
";
        // Bullet line is on line 3 (1-indexed)
        let mut pending = vec![PendingId {
            line: 3,
            id: "01JINV000000000000000000AA".to_string(),
            kind: WriteBackKind::InvolvedIn {
                entity_name: "John Doe".to_string(),
            },
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[2], "- John Doe");
        assert_eq!(lines[3], "  id: 01JINV000000000000000000AA");
    }

    #[test]
    fn writeback_involved_in_new_section() {
        let content = "\
---
id: 01CASE000000000000000000AA
sources:
  - https://example.com
---

# Some Case

Summary text.

## Events

### Something
- occurred_at: 2024-01-01

## Timeline

- Something -> Other thing
";
        let mut pending = vec![PendingId {
            line: 0,
            id: "01JINV000000000000000000BB".to_string(),
            kind: WriteBackKind::InvolvedIn {
                entity_name: "Alice".to_string(),
            },
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        assert!(result.contains("## Involved"));
        assert!(result.contains("- Alice"));
        assert!(result.contains("  id: 01JINV000000000000000000BB"));

        // ## Involved should appear before ## Timeline
        let involved_pos = result.find("## Involved").unwrap();
        let timeline_pos = result.find("## Timeline").unwrap();
        assert!(involved_pos < timeline_pos);
    }

    #[test]
    fn writeback_involved_in_new_section_multiple_entities() {
        let content = "\
---
sources:
  - https://example.com
---

# Case

Summary.

## Timeline

- A -> B
";
        let mut pending = vec![
            PendingId {
                line: 0,
                id: "01JINV000000000000000000CC".to_string(),
                kind: WriteBackKind::InvolvedIn {
                    entity_name: "Alice".to_string(),
                },
            },
            PendingId {
                line: 0,
                id: "01JINV000000000000000000DD".to_string(),
                kind: WriteBackKind::InvolvedIn {
                    entity_name: "Bob Corp".to_string(),
                },
            },
        ];

        let result = apply_writebacks(content, &mut pending).unwrap();
        assert!(result.contains("## Involved"));
        assert!(result.contains("- Alice"));
        assert!(result.contains("  id: 01JINV000000000000000000CC"));
        assert!(result.contains("- Bob Corp"));
        assert!(result.contains("  id: 01JINV000000000000000000DD"));

        // Section should be before ## Timeline
        let involved_pos = result.find("## Involved").unwrap();
        let timeline_pos = result.find("## Timeline").unwrap();
        assert!(involved_pos < timeline_pos);
    }

    #[test]
    fn writeback_timeline_edge() {
        let content = "\
## Timeline

- Event A -> Event B
";
        // Timeline bullet is on line 3 (1-indexed)
        let mut pending = vec![PendingId {
            line: 3,
            id: "01JTIM000000000000000000AA".to_string(),
            kind: WriteBackKind::TimelineEdge,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[2], "- Event A -> Event B");
        assert_eq!(lines[3], "  id: 01JTIM000000000000000000AA");
    }

    #[test]
    fn writeback_entity_front_matter_replaces_empty_id() {
        let content = "---\nid:\n---\n\n# Ali Murtopo\n\n- nationality: Indonesian\n";
        let end_line = find_front_matter_end(content).unwrap();
        assert_eq!(end_line, 3);

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JABC000000000000000000AA".to_string(),
            kind: WriteBackKind::EntityFrontMatter,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[0], "---");
        assert_eq!(lines[1], "id: 01JABC000000000000000000AA");
        assert_eq!(lines[2], "---");
        // Must NOT have a duplicate id line
        assert_eq!(lines.len(), 7);
    }

    #[test]
    fn writeback_case_id_replaces_empty_id() {
        let content = "---\nid:\nsources:\n  - https://example.com\n---\n\n# Some Case\n";
        let end_line = find_front_matter_end(content).unwrap();
        assert_eq!(end_line, 5);

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JXYZ123456789ABCDEFGHIJK".to_string(),
            kind: WriteBackKind::CaseId,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        assert_eq!(lines[1], "id: 01JXYZ123456789ABCDEFGHIJK");
        // No duplicate — same line count as original
        assert_eq!(lines.len(), 7);
    }

    #[test]
    fn writeback_does_not_replace_populated_id() {
        let content = "---\nid: 01JEXISTING00000000000000\n---\n\n# Test\n";
        let end_line = find_front_matter_end(content).unwrap();

        let mut pending = vec![PendingId {
            line: end_line,
            id: "01JNEW000000000000000000AA".to_string(),
            kind: WriteBackKind::EntityFrontMatter,
        }];

        let result = apply_writebacks(content, &mut pending).unwrap();
        let lines: Vec<&str> = result.lines().collect();
        // Should insert (not replace) since existing id is not empty
        assert_eq!(lines[1], "id: 01JEXISTING00000000000000");
        assert_eq!(lines[2], "id: 01JNEW000000000000000000AA");
    }
}