rbt_weave

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

impl WriteBackKind {
    /// The formatted `id: <NULID>` text to insert.
    fn format_id(&self, id: &str) -> String {
        match self {
            Self::EntityFrontMatter | Self::CaseId => format!("id: {id}"),
            Self::InlineEvent => format!("- id: {id}"),
            Self::Relationship
            | Self::RelatedCase
            | Self::InvolvedIn { .. }
            | Self::TimelineEdge => format!("  id: {id}"),
        }
    }

    /// Whether this kind targets YAML front matter (insert before closing `---`).
    const fn is_front_matter(&self) -> bool {
        matches!(self, Self::EntityFrontMatter | Self::CaseId)
    }

    /// Whether this kind creates a new section when `line == 0`.
    const fn needs_section_creation(&self) -> bool {
        matches!(self, Self::InvolvedIn { .. })
    }
}

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

    let mut lines: Vec<String> = content.lines().map(String::from).collect();
    let trailing_newline = content.ends_with('\n');

    // Partition: section-creation entries (InvolvedIn with line==0) vs normal
    let (new_involved, mut normal): (Vec<_>, Vec<_>) = pending
        .iter()
        .partition::<Vec<_>, _>(|p| p.kind.needs_section_creation() && p.line == 0);

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

    for p in &normal {
        let text = p.kind.format_id(&p.id);

        if p.kind.is_front_matter() {
            insert_front_matter_id(&mut lines, p.line, &text);
        } else {
            insert_body_id(&mut lines, p.line, &text);
        }
    }

    // Append ## Involved section for entities that need a new section
    if !new_involved.is_empty() {
        let entries: Vec<_> = new_involved
            .iter()
            .filter_map(|p| match &p.kind {
                WriteBackKind::InvolvedIn { entity_name } => Some((entity_name.as_str(), &*p.id)),
                _ => None,
            })
            .collect();
        append_involved_section(&mut lines, &entries);
    }

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

/// Insert an ID into YAML front matter, replacing an empty `id:` if present.
fn insert_front_matter_id(lines: &mut Vec<String>, closing_line: usize, text: &str) {
    let end_idx = closing_line.saturating_sub(1); // 0-indexed
    let bound = end_idx.min(lines.len());

    // Look for an existing empty `id:` line to replace
    for line in lines.iter_mut().take(bound) {
        let trimmed = line.trim();
        if trimmed == "id:" || trimmed == "id: " {
            *line = text.to_string();
            return;
        }
    }

    // No empty id found — insert before the closing `---`
    if end_idx <= lines.len() {
        lines.insert(end_idx, text.to_string());
    }
}

/// Insert an ID after a body element (event heading, relationship, etc.),
/// replacing an existing empty `id:` or skipping if already populated.
fn insert_body_id(lines: &mut Vec<String>, parent_line: usize, text: &str) {
    // Scan indented lines after the parent for an existing id field
    for line in lines.iter_mut().skip(parent_line) {
        let trimmed = line.trim();

        // Check for id field
        if let Some(value) = strip_id_prefix(trimmed) {
            if value.is_empty() {
                // Empty id — replace, preserving original indentation
                let indent = &line[..line.len() - line.trim_start().len()];
                *line = format!("{indent}{}", text.trim_start());
            }
            // Existing id (empty or populated) — don't insert another
            return;
        }

        // Stop at blank line, heading, or unindented line (next top-level bullet)
        if trimmed.is_empty() || trimmed.starts_with('#') || !line.starts_with(' ') {
            break;
        }
    }

    // No existing id found — insert after the parent line
    if parent_line <= lines.len() {
        lines.insert(parent_line, text.to_string());
    }
}

/// Strip `id:` prefix, returning the value portion (possibly empty).
/// Returns `None` if the line is not an id field.
fn strip_id_prefix(trimmed: &str) -> Option<&str> {
    trimmed.strip_prefix("id:").map(str::trim)
}

/// Append entries to an existing `## Involved` section, or create a new one.
/// Inserts before `## Timeline` or `## Related Cases`, or at end of file.
fn append_involved_section(lines: &mut Vec<String>, entries: &[(&str, &str)]) {
    // Check if ## Involved already exists
    if let Some(existing_idx) = lines.iter().position(|l| l.trim() == "## Involved") {
        // Find the end of the existing Involved section (next ## heading or EOF)
        let mut insert_at = existing_idx + 1;
        for (i, line) in lines.iter().enumerate().skip(existing_idx + 1) {
            if line.trim().starts_with("## ") {
                break;
            }
            insert_at = i + 1;
        }

        // Insert entries at the end of the existing section
        let mut offset = 0;
        for (name, id) in entries {
            lines.insert(insert_at + offset, format!("- {name}"));
            offset += 1;
            lines.insert(insert_at + offset, format!("  id: {id}"));
            offset += 1;
        }
    } else {
        // No existing section — create a new one
        let insert_idx = find_section_insert_point(lines);

        let mut section = Vec::with_capacity(2 + entries.len() * 2);

        // Blank line before section if needed
        if insert_idx > 0 && !lines[insert_idx - 1].is_empty() {
            section.push(String::new());
        }
        section.push("## Involved".to_string());
        section.push(String::new());

        for (name, id) in entries {
            section.push(format!("- {name}"));
            section.push(format!("  id: {id}"));
        }

        for (offset, line) in section.into_iter().enumerate() {
            lines.insert(insert_idx + offset, line);
        }
    }
}

/// 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 {
    lines
        .iter()
        .position(|l| {
            let t = l.trim();
            t == "## Timeline" || t == "## Related Cases"
        })
        .unwrap_or(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 1-indexed line number of the closing `---` in YAML front matter.
/// Returns `None` if the file has no valid front matter.
pub fn find_front_matter_end(content: &str) -> Option<usize> {
    let mut in_front_matter = false;
    for (i, line) in content.lines().enumerate() {
        if line.trim() == "---" {
            if in_front_matter {
                return Some(i + 1);
            }
            in_front_matter = true;
        }
    }
    None
}

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

    // -- Front matter insertion --

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

        let result = apply(
            &[pending(
                end_line,
                "01JXYZ",
                WriteBackKind::EntityFrontMatter,
            )],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[0], "---");
        assert_eq!(lines[1], "id: 01JXYZ");
        assert_eq!(lines[2], "---");
    }

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

        let result = apply(
            &[pending(
                end_line,
                "01JABC",
                WriteBackKind::EntityFrontMatter,
            )],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[1], "other: value");
        assert_eq!(lines[2], "id: 01JABC");
        assert_eq!(lines[3], "---");
    }

    #[test]
    fn 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();

        let result = apply(
            &[pending(
                end_line,
                "01JABC",
                WriteBackKind::EntityFrontMatter,
            )],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[1], "id: 01JABC");
        assert_eq!(lines[2], "---");
        assert_eq!(lines.len(), 7); // no duplicate
    }

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

        let result = apply(
            &[pending(end_line, "01JXYZ", WriteBackKind::CaseId)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[3], "id: 01JXYZ");
        assert_eq!(lines[4], "---");
    }

    #[test]
    fn 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();

        let result = apply(
            &[pending(end_line, "01JXYZ", WriteBackKind::CaseId)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[1], "id: 01JXYZ");
        assert_eq!(lines.len(), 7); // no duplicate
    }

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

        let result = apply(
            &[pending(
                end_line,
                "01JNEW",
                WriteBackKind::EntityFrontMatter,
            )],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[1], "id: 01JEXISTING");
        assert_eq!(lines[2], "id: 01JNEW"); // inserts, doesn't replace
    }

    // -- Inline event --

    #[test]
    fn inline_event() {
        let content =
            "## Events\n\n### Dismissal\n- occurred_at: 2024-12-24\n- event_type: termination\n";

        let result = apply(&[pending(3, "01JXYZ", WriteBackKind::InlineEvent)], content);
        let lines = to_lines(&result);
        assert_eq!(lines[2], "### Dismissal");
        assert_eq!(lines[3], "- id: 01JXYZ");
        assert_eq!(lines[4], "- occurred_at: 2024-12-24");
    }

    // -- Relationship --

    #[test]
    fn relationship_insert() {
        let content =
            "## Relationships\n\n- Alice -> Bob: employed_by\n  - source: https://example.com\n";

        let result = apply(
            &[pending(3, "01JXYZ", WriteBackKind::Relationship)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[2], "- Alice -> Bob: employed_by");
        assert_eq!(lines[3], "  id: 01JXYZ");
        assert_eq!(lines[4], "  - source: https://example.com");
    }

    #[test]
    fn relationship_replaces_empty_id() {
        let content = "## Relationships\n\n- A -> B: preceded_by\n  id:\n  description: replaced\n";

        let result = apply(
            &[pending(3, "01JABC", WriteBackKind::Relationship)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[3], "  id: 01JABC");
        assert_eq!(lines[4], "  description: replaced");
        assert_eq!(lines.len(), 5);
    }

    #[test]
    fn relationship_does_not_duplicate_populated_id() {
        let content =
            "## Relationships\n\n- A -> B: preceded_by\n  id: 01JEXISTING\n  description: test\n";

        let result = apply(
            &[pending(3, "01JNEW", WriteBackKind::Relationship)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[3], "  id: 01JEXISTING");
        assert_eq!(lines.len(), 5);
    }

    #[test]
    fn relationship_does_not_replace_old_bullet_format() {
        // `- id:` is NOT recognized anymore — inserts a new `  id:` line
        let content = "## Relationships\n\n- A -> B: preceded_by\n  - id:\n  description: test\n";

        let result = apply(
            &[pending(3, "01JABC", WriteBackKind::Relationship)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[3], "  id: 01JABC"); // new id inserted after parent
        assert_eq!(lines[4], "  - id:"); // old bullet format left as-is
    }

    // -- Related case --

    #[test]
    fn related_case() {
        let content =
            "## Related Cases\n\n- id/corruption/2013/some-case\n  description: Related scandal\n";

        let result = apply(&[pending(3, "01JREL", WriteBackKind::RelatedCase)], content);
        let lines = to_lines(&result);
        assert_eq!(lines[2], "- id/corruption/2013/some-case");
        assert_eq!(lines[3], "  id: 01JREL");
        assert_eq!(lines[4], "  description: Related scandal");
    }

    // -- Involved --

    #[test]
    fn involved_in_existing_section() {
        let content = "## Involved\n\n- John Doe\n";
        let kind = WriteBackKind::InvolvedIn {
            entity_name: "John Doe".to_string(),
        };

        let result = apply(&[pending(3, "01JINV", kind)], content);
        let lines = to_lines(&result);
        assert_eq!(lines[2], "- John Doe");
        assert_eq!(lines[3], "  id: 01JINV");
    }

    #[test]
    fn involved_in_new_section() {
        let content = "---\nid: 01CASE\nsources:\n  - https://example.com\n---\n\n# Some Case\n\nSummary.\n\n## Events\n\n### Something\n- occurred_at: 2024-01-01\n\n## Timeline\n\n- Something -> Other thing\n";
        let kind = WriteBackKind::InvolvedIn {
            entity_name: "Alice".to_string(),
        };

        let result = apply(&[pending(0, "01JINV", kind)], content);
        assert!(result.contains("## Involved"));
        assert!(result.contains("- Alice"));
        assert!(result.contains("  id: 01JINV"));

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

    #[test]
    fn involved_in_new_section_multiple_entities() {
        let content = "---\nsources:\n  - https://example.com\n---\n\n# Case\n\nSummary.\n\n## Timeline\n\n- A -> B\n";

        let result = apply(
            &[
                pending(
                    0,
                    "01JCC",
                    WriteBackKind::InvolvedIn {
                        entity_name: "Alice".to_string(),
                    },
                ),
                pending(
                    0,
                    "01JDD",
                    WriteBackKind::InvolvedIn {
                        entity_name: "Bob Corp".to_string(),
                    },
                ),
            ],
            content,
        );
        assert!(result.contains("- Alice"));
        assert!(result.contains("  id: 01JCC"));
        assert!(result.contains("- Bob Corp"));
        assert!(result.contains("  id: 01JDD"));

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

    #[test]
    fn involved_in_appends_to_existing_section() {
        let content = "## Involved\n\n- John Doe\n  id: 01JEXIST\n\n## Timeline\n\n- A -> B\n";
        let result = apply(
            &[pending(
                0,
                "01JNEW",
                WriteBackKind::InvolvedIn {
                    entity_name: "Event X".to_string(),
                },
            )],
            content,
        );
        let lines = to_lines(&result);
        // Should have only one ## Involved section
        let involved_count = lines.iter().filter(|l| l.trim() == "## Involved").count();
        assert_eq!(
            involved_count, 1,
            "should not create duplicate ## Involved section"
        );
        // The new entry should be in the section
        assert!(result.contains("- Event X"));
        assert!(result.contains("  id: 01JNEW"));
        // Original entry preserved
        assert!(result.contains("- John Doe"));
        assert!(result.contains("  id: 01JEXIST"));
    }

    // -- Timeline edge --

    #[test]
    fn timeline_edge() {
        let content = "## Timeline\n\n- Event A -> Event B\n";

        let result = apply(
            &[pending(3, "01JTIM", WriteBackKind::TimelineEdge)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[2], "- Event A -> Event B");
        assert_eq!(lines[3], "  id: 01JTIM");
    }

    #[test]
    fn timeline_edge_replaces_empty_id() {
        let content = "## Timeline\n\n- Event A -> Event B\n  id:\n";

        let result = apply(
            &[pending(3, "01JABC", WriteBackKind::TimelineEdge)],
            content,
        );
        let lines = to_lines(&result);
        assert_eq!(lines[3], "  id: 01JABC");
        assert_eq!(lines.len(), 4);
    }

    // -- Multiple insertions --

    #[test]
    fn multiple_insertions() {
        let content = "## Events\n\n### Event A\n- occurred_at: 2024-01-01\n\n### Event B\n- occurred_at: 2024-06-01\n\n## Relationships\n\n- Event A -> Event B: associate_of\n";

        let result = apply(
            &[
                pending(3, "01JAAA", WriteBackKind::InlineEvent),
                pending(6, "01JBBB", WriteBackKind::InlineEvent),
                pending(10, "01JCCC", WriteBackKind::Relationship),
            ],
            content,
        );
        assert!(result.contains("- id: 01JAAA"));
        assert!(result.contains("- id: 01JBBB"));
        assert!(result.contains("  id: 01JCCC"));
    }

    // -- Misc --

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

    #[test]
    fn preserves_trailing_newline() {
        let content = "---\n---\n\n# Test\n";
        let result = apply(
            &[pending(2, "01JABC", WriteBackKind::EntityFrontMatter)],
            content,
        );
        assert!(result.ends_with('\n'));
    }

    #[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);
    }

    // -- Helpers --

    fn pending(line: usize, id: &str, kind: WriteBackKind) -> PendingId {
        PendingId {
            line,
            id: id.to_string(),
            kind,
        }
    }

    fn apply(specs: &[PendingId], content: &str) -> String {
        // apply_writebacks takes &mut [PendingId], so we need an owned vec
        let mut owned: Vec<PendingId> = specs
            .iter()
            .map(|p| PendingId {
                line: p.line,
                id: p.id.clone(),
                kind: match &p.kind {
                    WriteBackKind::EntityFrontMatter => WriteBackKind::EntityFrontMatter,
                    WriteBackKind::CaseId => WriteBackKind::CaseId,
                    WriteBackKind::InlineEvent => WriteBackKind::InlineEvent,
                    WriteBackKind::Relationship => WriteBackKind::Relationship,
                    WriteBackKind::RelatedCase => WriteBackKind::RelatedCase,
                    WriteBackKind::InvolvedIn { entity_name } => WriteBackKind::InvolvedIn {
                        entity_name: entity_name.clone(),
                    },
                    WriteBackKind::TimelineEdge => WriteBackKind::TimelineEdge,
                },
            })
            .collect();
        apply_writebacks(content, &mut owned).unwrap()
    }

    fn to_lines(s: &str) -> Vec<&str> {
        s.lines().collect()
    }
}