extend-docx 0.1.0

use std::collections::{HashMap, HashSet};

use std::cell::RefCell;

use crate::model::{
    DocNode, DocumentCommentDefinition, DocumentNoteDefinition, DocumentSection, FooterSection,
    FormFieldRunNode, FormFieldType, HeaderSection, ParagraphChildNode, ParagraphNode,
    TableCellContentNode,
};
use crate::package::OoxmlPackage;
use crate::parse::re;
use crate::xml::{
    decode_xml_entities, extract_balanced_tag_blocks, get_attribute, parse_integer_attribute,
};

use super::body::body_xml;
use super::context::{ContentTypeLookup, ParseContext, ParsedStyleSheet};
use super::document::{
    extract_body_token_ranges, is_go_back_bookmark_paragraph, parse_document_xml, BodyTokenKind,
};
use super::metadata::extract_section_properties_xml;
use super::relationships::parse_part_relationships;
use super::util::on_off_value_to_boolean;

struct ReferencedSectionPart {
    part_name: String,
    reference_type: Option<String>,
    nodes: Vec<DocNode>,
}

pub(crate) struct PartNodeCache<'a, 's> {
    pkg: &'a OoxmlPackage,
    binary_assets: &'a HashMap<String, Vec<u8>>,
    content_types: &'s ContentTypeLookup,
    style_sheet: &'s ParsedStyleSheet,
    parsed_part_nodes_by_name: HashMap<String, Vec<DocNode>>,
    context: Option<ParseContext<'a, &'s ParsedStyleSheet>>,
}

impl<'a, 's> PartNodeCache<'a, 's> {
    pub(crate) fn new(
        pkg: &'a OoxmlPackage,
        content_types: &'s ContentTypeLookup,
        style_sheet: &'s ParsedStyleSheet,
        binary_assets: &'a HashMap<String, Vec<u8>>,
    ) -> Self {
        Self {
            pkg,
            binary_assets,
            content_types,
            style_sheet,
            parsed_part_nodes_by_name: HashMap::new(),
            context: None,
        }
    }

    fn resolve(
        &mut self,
        part_name: &str,
        relationship_tag_name: &str,
        warnings: &mut Vec<String>,
    ) -> Option<Vec<DocNode>> {
        if let Some(cached_nodes) = self.parsed_part_nodes_by_name.get(part_name) {
            return Some(cached_nodes.clone());
        }

        let Some(part_xml) = self
            .pkg
            .parts
            .get(part_name)
            .map(|part| part.content.as_str())
        else {
            warnings.push(format!(
                "{} {}",
                if relationship_tag_name == "headerReference" {
                    "Missing header part"
                } else {
                    "Missing footer part"
                },
                part_name
            ));
            return None;
        };

        if self.context.is_none() {
            self.context = Some(ParseContext {
                relationships: HashMap::new(),
                content_types: self.content_types.clone(),
                parts: &self.pkg.parts,
                binary_assets: self.binary_assets,
                style_sheet: self.style_sheet,
                warnings: RefCell::new(Vec::new()),
            });
        }
        let context = self.context.as_mut().unwrap();
        context.relationships = parse_part_relationships(self.pkg, part_name);
        context.warnings.get_mut().clear();
        let nodes = parse_document_xml(part_xml, context);
        self.parsed_part_nodes_by_name
            .insert(part_name.to_string(), nodes.clone());
        Some(nodes)
    }
}

pub fn parse_header_sections(
    pkg: &OoxmlPackage,
    document_xml: &str,
    document_relationships: &HashMap<String, String>,
    warnings: &mut Vec<String>,
    cache: &mut PartNodeCache<'_, '_>,
) -> Vec<HeaderSection> {
    parse_referenced_sections(
        pkg,
        document_xml,
        document_relationships,
        warnings,
        cache,
        "headerReference",
        "Missing header relationship target for",
        "Missing header part",
    )
    .into_iter()
    .map(|section| HeaderSection {
        part_name: section.part_name,
        reference_type: section.reference_type,
        nodes: section.nodes,
    })
    .collect()
}

pub fn parse_footer_sections(
    pkg: &OoxmlPackage,
    document_xml: &str,
    document_relationships: &HashMap<String, String>,
    warnings: &mut Vec<String>,
    cache: &mut PartNodeCache<'_, '_>,
) -> Vec<FooterSection> {
    parse_referenced_sections(
        pkg,
        document_xml,
        document_relationships,
        warnings,
        cache,
        "footerReference",
        "Missing footer relationship target for",
        "Missing footer part",
    )
    .into_iter()
    .map(|section| FooterSection {
        part_name: section.part_name,
        reference_type: section.reference_type,
        nodes: section.nodes,
    })
    .collect()
}

pub fn parse_document_sections(
    document_xml: &str,
    document_relationships: &HashMap<String, String>,
    warnings: &mut Vec<String>,
    cache: &mut PartNodeCache<'_, '_>,
) -> Vec<DocumentSection> {
    let body_xml = body_xml(document_xml);
    let token_ranges = extract_body_token_ranges(&body_xml);
    let mut node_count: i64 = 0;
    let mut section_start_node_index: i64 = 0;
    let mut sections = Vec::new();

    for token in token_ranges {
        let token_xml = &body_xml[token.start..token.end];
        let produced_node =
            token.kind == BodyTokenKind::Table || !is_go_back_bookmark_paragraph(token_xml);

        if token.kind == BodyTokenKind::Paragraph {
            if let Some(section_properties_xml) = extract_balanced_tag_blocks(token_xml, "w:sectPr")
                .into_iter()
                .next()
            {
                sections.push(DocumentSection {
                    start_node_index: section_start_node_index,
                    section_properties_xml: Some(section_properties_xml.clone()),
                    header_sections: parse_section_references_from_properties(
                        &section_properties_xml,
                        "headerReference",
                        document_relationships,
                        cache,
                        warnings,
                    )
                    .into_iter()
                    .map(|section| HeaderSection {
                        part_name: section.part_name,
                        reference_type: section.reference_type,
                        nodes: section.nodes,
                    })
                    .collect(),
                    footer_sections: parse_section_references_from_properties(
                        &section_properties_xml,
                        "footerReference",
                        document_relationships,
                        cache,
                        warnings,
                    )
                    .into_iter()
                    .map(|section| FooterSection {
                        part_name: section.part_name,
                        reference_type: section.reference_type,
                        nodes: section.nodes,
                    })
                    .collect(),
                });
                section_start_node_index = node_count + if produced_node { 1 } else { 0 };
            }
        }

        if produced_node {
            node_count += 1;
        }
    }

    let final_section_properties_xml = extract_section_properties_xml(document_xml);
    if let Some(section_properties_xml) = final_section_properties_xml.clone() {
        sections.push(DocumentSection {
            start_node_index: section_start_node_index,
            section_properties_xml: Some(section_properties_xml.clone()),
            header_sections: parse_section_references_from_properties(
                &section_properties_xml,
                "headerReference",
                document_relationships,
                cache,
                warnings,
            )
            .into_iter()
            .map(|section| HeaderSection {
                part_name: section.part_name,
                reference_type: section.reference_type,
                nodes: section.nodes,
            })
            .collect(),
            footer_sections: parse_section_references_from_properties(
                &section_properties_xml,
                "footerReference",
                document_relationships,
                cache,
                warnings,
            )
            .into_iter()
            .map(|section| FooterSection {
                part_name: section.part_name,
                reference_type: section.reference_type,
                nodes: section.nodes,
            })
            .collect(),
        });
    }

    let mut normalized_sections = Vec::new();
    for section in sections {
        if node_count > 0 && section.start_node_index > node_count {
            continue;
        }

        let duplicate = normalized_sections
            .last()
            .is_some_and(|previous: &DocumentSection| {
                previous.start_node_index == section.start_node_index
                    && previous.section_properties_xml == section.section_properties_xml
            });
        if duplicate {
            continue;
        }
        normalized_sections.push(section);
    }

    if normalized_sections.is_empty() {
        normalized_sections.push(DocumentSection {
            start_node_index: 0,
            section_properties_xml: final_section_properties_xml,
            header_sections: Vec::new(),
            footer_sections: Vec::new(),
        });
    }

    normalized_sections.sort_by_key(|section| section.start_node_index);
    normalized_sections
}

pub fn parse_document_notes_from_part(
    pkg: &OoxmlPackage,
    part_name: &str,
    tag_name: &str,
    content_types: &ContentTypeLookup,
    style_sheet: &ParsedStyleSheet,
    warnings: &mut Vec<String>,
    binary_assets: &HashMap<String, Vec<u8>>,
) -> Vec<DocumentNoteDefinition> {
    let _ = warnings;
    let Some(notes_xml) = pkg.parts.get(part_name).map(|part| part.content.as_str()) else {
        return Vec::new();
    };

    let context = ParseContext {
        relationships: parse_part_relationships(pkg, part_name),
        content_types: content_types.clone(),
        parts: &pkg.parts,
        binary_assets,
        style_sheet,
        warnings: RefCell::new(Vec::new()),
    };

    let mut notes = Vec::new();
    for note_xml in crate::xml::extract_balanced_tag_slices(notes_xml, tag_name) {
        let note_tag = find_opening_tag(&note_xml, tag_name).unwrap_or_default();
        if get_attribute(&note_tag, "w:type").is_some() {
            continue;
        }

        let note_id = parse_integer_attribute(&note_tag, "w:id");
        let Some(note_id) = note_id.filter(|id| *id >= 0) else {
            continue;
        };

        let parsed_nodes = parse_document_xml(&note_xml, &context);
        let text = note_text_from_nodes(&parsed_nodes);
        if text.is_empty() {
            continue;
        }

        notes.push(DocumentNoteDefinition {
            id: note_id,
            text,
            nodes: Some(parsed_nodes),
        });
    }

    notes.sort_by_key(|note| note.id);
    notes
}

/// Parses `word/comments.xml` into comment definitions, enriched with
/// resolution state and threading from `word/commentsExtended.xml` (matched
/// through the `w14:paraId` of each comment's final paragraph, per MS-DOCX).
pub fn parse_document_comments(
    pkg: &OoxmlPackage,
    content_types: &ContentTypeLookup,
    style_sheet: &ParsedStyleSheet,
    binary_assets: &HashMap<String, Vec<u8>>,
) -> Vec<DocumentCommentDefinition> {
    let Some(comments_xml) =
        crate::package::related_document_part(pkg, "comments").map(|part| part.content.as_str())
    else {
        return Vec::new();
    };

    let context = ParseContext {
        relationships: parse_part_relationships(
            pkg,
            &crate::package::related_document_part_name(pkg, "comments"),
        ),
        content_types: content_types.clone(),
        parts: &pkg.parts,
        binary_assets,
        style_sheet,
        warnings: RefCell::new(Vec::new()),
    };

    // commentsExtended: paraId -> (done, parent paraId).
    let mut extended_by_para_id: HashMap<String, (Option<bool>, Option<String>)> = HashMap::new();
    if let Some(extended_xml) = pkg
        .parts
        .get("word/commentsExtended.xml")
        .map(|part| part.content.as_str())
    {
        for tag in re::get_unchecked(r"(?i)<w15:commentEx\b[^>]*/?>")
            .find_iter(extended_xml)
            .map(|m| m.as_str())
        {
            let Some(para_id) = get_attribute(tag, "w15:paraId") else {
                continue;
            };
            let done_attribute = get_attribute(tag, "w15:done");
            let done = on_off_value_to_boolean(done_attribute.as_deref());
            let parent_para_id = get_attribute(tag, "w15:paraIdParent");
            extended_by_para_id.insert(para_id, (done, parent_para_id));
        }
    }

    let mut comments = Vec::new();
    let mut comment_id_by_para_id: HashMap<String, i64> = HashMap::new();
    let mut parent_para_id_by_comment_id: HashMap<i64, String> = HashMap::new();
    for comment_xml in extract_balanced_tag_blocks(comments_xml, "w:comment") {
        let comment_tag = find_opening_tag(&comment_xml, "w:comment").unwrap_or_default();
        let Some(comment_id) = parse_integer_attribute(&comment_tag, "w:id") else {
            continue;
        };

        let parsed_nodes = parse_document_xml(&comment_xml, &context);
        let text = comment_text_from_nodes(&parsed_nodes);
        let author = get_attribute(&comment_tag, "w:author")
            .map(|value| decode_xml_entities(&value))
            .filter(|value| !value.trim().is_empty());
        let initials = get_attribute(&comment_tag, "w:initials")
            .map(|value| decode_xml_entities(&value))
            .filter(|value| !value.trim().is_empty());
        let date = get_attribute(&comment_tag, "w:date").filter(|value| !value.trim().is_empty());

        // The comment's paragraphs carry w14:paraId attributes; the LAST one
        // identifies the comment in commentsExtended.
        let para_ids: Vec<String> = extract_balanced_tag_blocks(&comment_xml, "w:p")
            .into_iter()
            .filter_map(|paragraph_xml| find_opening_tag(&paragraph_xml, "w:p"))
            .filter_map(|paragraph_tag| get_attribute(&paragraph_tag, "w14:paraId"))
            .collect();
        let extended_paragraph_id = para_ids.last().cloned();
        let mut resolved = None;
        if let Some(last_para_id) = extended_paragraph_id.as_ref() {
            for para_id in &para_ids {
                comment_id_by_para_id.insert(para_id.clone(), comment_id);
            }
            if let Some((done, parent_para_id)) = extended_by_para_id.get(last_para_id) {
                // `done="0"` is meaningful provenance. Do not collapse an
                // explicitly unresolved thread into a missing value.
                resolved = *done;
                if let Some(parent_para_id) = parent_para_id {
                    parent_para_id_by_comment_id.insert(comment_id, parent_para_id.clone());
                }
            }
        }

        comments.push(DocumentCommentDefinition {
            nodes: Some(parsed_nodes),
            id: comment_id,
            author,
            initials,
            date,
            text,
            parent_id: None,
            resolved,
            source_xml: Some(comment_xml),
            extended_paragraph_id,
            source_resolved: resolved,
            resolution_dirty: Some(false),
            is_new: Some(false),
        });
    }

    for comment in comments.iter_mut() {
        if let Some(parent_para_id) = parent_para_id_by_comment_id.get(&comment.id) {
            comment.parent_id = comment_id_by_para_id.get(parent_para_id).copied();
        }
    }

    comments.sort_by_key(|comment| comment.id);
    comments
}

fn parse_referenced_sections(
    pkg: &OoxmlPackage,
    document_xml: &str,
    document_relationships: &HashMap<String, String>,
    warnings: &mut Vec<String>,
    cache: &mut PartNodeCache<'_, '_>,
    relationship_tag_name: &str,
    missing_target_prefix: &str,
    missing_part_prefix: &str,
) -> Vec<ReferencedSectionPart> {
    let mut sections = Vec::new();
    let mut seen_part_names = HashSet::new();

    for tag in find_opening_tags(document_xml, relationship_tag_name) {
        let relationship_id = get_attribute(&tag, "r:id");
        let Some(relationship_id) = relationship_id else {
            continue;
        };

        let Some(target_part_name) = document_relationships.get(&relationship_id) else {
            warnings.push(format!("{missing_target_prefix} {relationship_id}"));
            continue;
        };

        if !seen_part_names.insert(target_part_name.clone()) {
            continue;
        }

        let Some(_header_xml) = pkg
            .parts
            .get(target_part_name)
            .map(|part| part.content.as_str())
        else {
            warnings.push(format!("{missing_part_prefix} {target_part_name}"));
            continue;
        };

        let Some(nodes) = cache.resolve(target_part_name, relationship_tag_name, warnings) else {
            continue;
        };

        sections.push(ReferencedSectionPart {
            part_name: target_part_name.clone(),
            reference_type: get_attribute(&tag, "w:type"),
            nodes,
        });
    }

    sections
}

fn parse_section_references_from_properties(
    section_properties_xml: &str,
    relationship_tag_name: &str,
    document_relationships: &HashMap<String, String>,
    cache: &mut PartNodeCache<'_, '_>,
    warnings: &mut Vec<String>,
) -> Vec<ReferencedSectionPart> {
    let mut references = Vec::new();
    let mut seen_references = HashSet::new();

    for tag in find_opening_tags(section_properties_xml, relationship_tag_name) {
        let relationship_id = get_attribute(&tag, "r:id");
        let Some(relationship_id) = relationship_id else {
            continue;
        };

        let Some(target_part_name) = document_relationships.get(&relationship_id) else {
            warnings.push(format!(
                "{} {}",
                if relationship_tag_name == "headerReference" {
                    "Missing header relationship target for"
                } else {
                    "Missing footer relationship target for"
                },
                relationship_id
            ));
            continue;
        };

        let reference_type = get_attribute(&tag, "w:type");
        let dedupe_key = format!(
            "{target_part_name}::{}",
            reference_type.as_deref().unwrap_or("")
        );
        if !seen_references.insert(dedupe_key) {
            continue;
        }

        let Some(nodes) = cache.resolve(target_part_name, relationship_tag_name, warnings) else {
            continue;
        };

        references.push(ReferencedSectionPart {
            part_name: target_part_name.clone(),
            reference_type,
            nodes,
        });
    }

    references
}

fn note_text_from_nodes(nodes: &[DocNode]) -> String {
    let mut lines = Vec::new();

    for node in nodes {
        append_paragraph_from_node(node, &mut lines);
    }

    let joined = lines.join("\n").trim().to_string();
    strip_leading_bracketed_index(&joined)
}

fn comment_text_from_nodes(nodes: &[DocNode]) -> String {
    let mut lines = Vec::new();
    for node in nodes {
        append_comment_text_from_node(node, &mut lines);
    }
    lines.join("\n")
}

fn append_comment_text_from_node(node: &DocNode, lines: &mut Vec<String>) {
    match node {
        DocNode::Paragraph(paragraph) => lines.push(paragraph_plain_text(paragraph)),
        DocNode::Table(table) => {
            for row in &table.rows {
                for cell in &row.cells {
                    for paragraph in cell_paragraphs_from_content(&cell.nodes) {
                        lines.push(paragraph_plain_text(&paragraph));
                    }
                }
            }
        }
    }
}

fn append_paragraph_from_node(node: &DocNode, lines: &mut Vec<String>) {
    match node {
        DocNode::Paragraph(paragraph) => append_paragraph(paragraph, lines),
        DocNode::Table(table) => {
            for row in &table.rows {
                for cell in &row.cells {
                    for paragraph in cell_paragraphs_from_content(&cell.nodes) {
                        append_paragraph(&paragraph, lines);
                    }
                }
            }
        }
    }
}

fn append_paragraph(paragraph: &ParagraphNode, lines: &mut Vec<String>) {
    let text = paragraph_plain_text(paragraph).trim_end().to_string();
    if !text.is_empty() {
        lines.push(text);
    }
}

fn paragraph_plain_text(paragraph: &ParagraphNode) -> String {
    paragraph
        .children
        .iter()
        .map(paragraph_child_plain_text)
        .collect::<Vec<_>>()
        .join("")
}

fn paragraph_child_plain_text(child: &ParagraphChildNode) -> String {
    match child {
        ParagraphChildNode::Text(text) => text.text.clone(),
        ParagraphChildNode::FormField(field) => form_field_plain_text(field),
        ParagraphChildNode::Image(_) => String::new(),
    }
}

fn form_field_plain_text(field: &FormFieldRunNode) -> String {
    match field.field_type {
        FormFieldType::Checkbox => {
            let is_checked = field.checked.unwrap_or_else(|| {
                field
                    .widget
                    .as_ref()
                    .and_then(|widget| widget.checkbox.as_ref())
                    .and_then(|checkbox| checkbox.default_checked)
                    .unwrap_or(false)
            });
            if is_checked {
                field
                    .checked_symbol
                    .clone()
                    .unwrap_or_else(|| "☒".to_string())
            } else {
                field
                    .unchecked_symbol
                    .clone()
                    .unwrap_or_else(|| "☐".to_string())
            }
        }
        FormFieldType::Text => field
            .value
            .clone()
            .or_else(|| {
                field
                    .widget
                    .as_ref()
                    .and_then(|widget| widget.text.as_ref())
                    .and_then(|text| text.default_text.clone())
            })
            .unwrap_or_default(),
        _ => field.value.clone().unwrap_or_default(),
    }
}

fn cell_paragraphs_from_content(nodes: &[TableCellContentNode]) -> Vec<ParagraphNode> {
    let mut paragraphs = Vec::new();
    walk_cell_content(nodes, &mut paragraphs);
    paragraphs
}

fn walk_cell_content(nodes: &[TableCellContentNode], paragraphs: &mut Vec<ParagraphNode>) {
    for item in nodes {
        match item {
            TableCellContentNode::Paragraph(paragraph) => paragraphs.push(paragraph.clone()),
            TableCellContentNode::Table(table) => {
                for row in &table.rows {
                    for cell in &row.cells {
                        walk_cell_content(&cell.nodes, paragraphs);
                    }
                }
            }
        }
    }
}

fn strip_leading_bracketed_index(text: &str) -> String {
    let bytes = text.as_bytes();
    if bytes.first() != Some(&b'[') {
        return text.to_string();
    }

    let mut index = 1;
    while index < bytes.len() && bytes[index].is_ascii_digit() {
        index += 1;
    }

    if index < bytes.len() && bytes[index] == b']' {
        let remainder = text[index + 1..].trim_start();
        return remainder.to_string();
    }

    text.to_string()
}

fn find_opening_tags<'a>(xml: &'a str, local_name: &str) -> Vec<&'a str> {
    let name = format!("w:{local_name}");
    if let Some(tags) = crate::xml_index::opening_tags(xml, &name) {
        return tags;
    }
    let open_prefix = format!("<{name}");
    let mut tags = Vec::new();
    let mut index = 0;

    while index < xml.len() {
        let Some(tag_start) = find_case_insensitive_at(xml, index, &open_prefix) else {
            break;
        };

        if !has_word_boundary_after(xml, tag_start + open_prefix.len()) {
            index = tag_start + 1;
            continue;
        }

        let Some(relative_gt) = xml[tag_start..].find('>') else {
            break;
        };

        let tag_end = tag_start + relative_gt + 1;
        tags.push(&xml[tag_start..tag_end]);
        index = tag_end;
    }

    tags
}

fn find_opening_tag(xml: &str, tag_name: &str) -> Option<String> {
    let open_prefix = format!("<{tag_name}");
    let tag_start = find_case_insensitive_at(xml, 0, &open_prefix)?;
    let relative_gt = xml[tag_start..].find('>')?;
    Some(xml[tag_start..tag_start + relative_gt + 1].to_string())
}

fn has_word_boundary_after(xml: &str, index: usize) -> bool {
    xml.as_bytes()
        .get(index)
        .is_none_or(|byte| !byte.is_ascii_alphanumeric() && *byte != b'_')
}

fn find_case_insensitive_at(haystack: &str, start: usize, needle: &str) -> Option<usize> {
    crate::xml::find_ignore_ascii_case(haystack, start, needle)
}

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

    #[test]
    fn indexed_openings_preserve_case_order_nesting_and_slice_boundaries() {
        let xml = r#"<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:headerReference w:type="default"><w:headerReference/></w:headerReference><w:HEADERREFERENCE/><w:headerReference_extra/><w:footerReference/></w:body></w:document>"#;
        let start = xml.find("<w:body>").unwrap();
        let end = xml.find("</w:body>").unwrap();
        let fragments = [xml, &xml[start..end], &xml[start + 1..end], &xml[..end - 2]];
        let expected: Vec<_> = fragments.iter().map(|fragment| find_opening_tags(fragment, "headerReference")).collect();
        let index = crate::xml_index::Index::canonical(xml).unwrap();
        let _scope = crate::xml_index::scope([("part", xml)].into_iter(), HashMap::from([("part".to_owned(), index)]));
        for (fragment, expected) in fragments.into_iter().zip(expected) {
            assert_eq!(find_opening_tags(fragment, "headerReference"), expected);
            assert!(find_opening_tags(fragment, "missingReference").is_empty());
        }
    }

    #[test]
    fn unindexed_openings_preserve_legacy_token_boundaries() {
        let xml = r#"<w:root><w:headerReference:extra/><w:headerReference-extra/><w:headerReference w:type="a>b"/></w:root>"#;
        let expected = find_opening_tags(xml, "headerReference");
        let _scope = crate::xml_index::scope([("part", xml)].into_iter(), HashMap::new());
        assert_eq!(find_opening_tags(xml, "headerReference"), expected);
        assert_eq!(expected.len(), 3);
    }
}