betteroffice-docx 0.3.0

use std::collections::BTreeMap;
use std::sync::Arc;

use docx_edit::content_controls::{
    ContentControlQuery, ContentControlsOptions, ContentControlsSnapshot,
    find_package_content_controls, list_package_content_controls,
};
use docx_edit::structured::{
    DocxStructuredContent, ExportOptions, MarkdownContent, MarkdownOptions,
    export_package_structured,
};
use docx_edit::{EditCtx, EditingDoc, Receipt, StoryRange};
use docx_layout::types::Input as LayoutInput;
use docx_parse::block::BlockContent;
use docx_parse::document::{DocumentBody, Section, get_paragraph_text};
use docx_parse::inline::{InlineNode, Run, RunContent, RunType};
use docx_parse::paragraph::{Paragraph, ParagraphContent};
use docx_parse::paragraph_identity::{
    allocate_paragraph_id, format_paragraph_id, package_paragraph_ids,
};
use docx_parse::s9::{
    S9DocumentBodyWire, S9DocumentWire, S9PackageWire, S9ParseOptions, S9SectionWire,
    S9WireEnvelope,
};
use docx_parse::serializer::{
    S13SaveOptions, S13SaveRequest, SerializerDeterminism, write_docx_s13_parts,
};
use docx_parse::table::Table;
use docx_parse::xml::ParseLimits;
use sha2::{Digest, Sha256};

use crate::types::DEFAULT_SERIALIZATION_TIME;
use crate::{DocumentModel, DocumentStructure, Error, LayoutResult, Result, SaveOptions};

pub struct Document {
    /// Inflated package parts retained from open for `save` to reuse.
    original_parts: Vec<(String, Vec<u8>)>,
    seed: String,
    model: DocumentModel,
    /// By body paragraph ordinal, the fresh ID handed to a paragraph that
    /// repeats an earlier paragraph's ID and the authored ID it saves with
    /// until it is edited or the model is changed.
    authored_ids: BTreeMap<usize, (String, String)>,
    #[cfg(feature = "raster")]
    pub(crate) fonts: crate::render::FontRegistry,
    #[cfg(feature = "raster")]
    pub(crate) images: crate::render::ImageRegistry,
}

impl Document {
    pub fn open(bytes: &[u8]) -> Result<Self> {
        Self::open_with_limits(bytes, &ParseLimits::default())
    }

    /// [`Document::open`] under a caller-supplied parse budget. A document that
    /// exceeds any limit is rejected rather than truncated.
    pub fn open_with_limits(bytes: &[u8], limits: &ParseLimits) -> Result<Self> {
        let (parsed, original_parts) = docx_parse::parse_docx_s9_wire_parts_with_limits(
            bytes,
            S9ParseOptions::default(),
            limits,
        )?;
        let mut document = parsed.document;
        let authored_ids =
            address_repeated_paragraphs(&mut document.package.document.content, &original_parts);
        let model = model_from_package(
            document.package,
            document.template_variables.unwrap_or_default(),
            document.warnings.unwrap_or_default(),
        );
        Ok(Self {
            original_parts,
            seed: format!("{:x}", Sha256::digest(bytes)),
            model,
            authored_ids,
            #[cfg(feature = "raster")]
            fonts: crate::render::FontRegistry::default(),
            #[cfg(feature = "raster")]
            images: crate::render::ImageRegistry::default(),
        })
    }

    pub fn model(&self) -> &DocumentModel {
        &self.model
    }

    /// The model, for any change. Paragraphs that repeat an earlier
    /// paragraph's ID then save with the fresh IDs they read with.
    pub fn model_mut(&mut self) -> &mut DocumentModel {
        self.authored_ids.clear();
        &mut self.model
    }

    pub fn into_model(self) -> DocumentModel {
        self.model
    }

    pub fn body(&self) -> &DocumentBody {
        &self.model.body
    }

    pub fn headers(&self) -> &[(String, docx_parse::HeaderFooter)] {
        &self.model.headers
    }

    pub fn footers(&self) -> &[(String, docx_parse::HeaderFooter)] {
        &self.model.footers
    }

    pub fn sections(&self) -> &[Section] {
        self.model.body.sections.as_deref().unwrap_or_default()
    }

    /// Every body paragraph in document order, table cells and content controls
    /// included. One that repeats an earlier paragraph's ID carries a fresh ID,
    /// which a save writes once the paragraph is edited.
    pub fn paragraphs(&self) -> Vec<&Paragraph> {
        let mut paragraphs = Vec::new();
        collect_paragraphs(&self.model.body.content, &mut paragraphs);
        paragraphs
    }

    pub fn tables(&self) -> Vec<&Table> {
        let mut tables = Vec::new();
        collect_tables(&self.model.body.content, &mut tables);
        tables
    }

    /// The body paragraph with `para_id`; `None` when no paragraph or several have it.
    pub fn paragraph(&self, para_id: &str) -> Option<&Paragraph> {
        self.unique_paragraph(para_id).ok()
    }

    fn unique_paragraph(&self, para_id: &str) -> Result<&Paragraph> {
        let mut matches = self
            .paragraphs()
            .into_iter()
            .filter(|paragraph| paragraph.para_id.as_deref() == Some(para_id));
        let found = matches
            .next()
            .ok_or_else(|| Error::ParagraphNotFound(para_id.to_owned()))?;
        match matches.next() {
            Some(_) => Err(Error::AmbiguousParagraph(para_id.to_owned())),
            None => Ok(found),
        }
    }

    pub fn structure(&self) -> DocumentStructure {
        DocumentStructure {
            body_paragraphs: self.paragraphs().len(),
            body_tables: self.tables().len(),
            sections: self.sections().len(),
            headers: self.model.headers.len(),
            footers: self.model.footers.len(),
            footnotes: self.model.footnotes.len(),
            endnotes: self.model.endnotes.len(),
        }
    }

    pub fn replace_paragraph_text(&mut self, para_id: &str, text: &str) -> Result<Receipt> {
        let context = EditCtx::local("betteroffice-docx", DEFAULT_SERIALIZATION_TIME);
        self.replace_paragraph_text_with(para_id, text, 1, &context)
    }

    pub fn replace_paragraph_text_with(
        &mut self,
        para_id: &str,
        text: &str,
        client_id: u64,
        context: &EditCtx,
    ) -> Result<Receipt> {
        self.unique_paragraph(para_id)?;
        let paragraph = find_paragraph_mut(&mut self.model.body.content, para_id)
            .ok_or_else(|| Error::ParagraphNotFound(para_id.to_owned()))?;
        let template = plain_run_template(paragraph)
            .ok_or_else(|| Error::UnsupportedParagraphEdit(para_id.to_owned()))?;
        let current = get_paragraph_text(paragraph);
        let style = paragraph
            .formatting
            .as_ref()
            .and_then(|formatting| formatting.style_id.as_deref())
            .unwrap_or("Normal");
        let alignment = paragraph
            .formatting
            .as_ref()
            .and_then(|formatting| formatting.alignment.as_deref())
            .unwrap_or("left");
        let editor = EditingDoc::new(client_id);
        editor.create_story_with_paragraph_id("body", para_id, &current, style, alignment)?;
        let receipt = editor.replace_range(
            context,
            StoryRange::new("body", 0, utf16_len(&current)),
            text,
        )?;
        let edited = editor
            .paragraphs("body")?
            .into_iter()
            .next()
            .map(|paragraph| paragraph.text)
            .unwrap_or_default();
        replace_plain_run(paragraph, edited, template);
        let replacement = paragraph.clone();
        if let Some(sections) = &mut self.model.body.sections {
            for section in sections {
                if let Some(section_paragraph) = find_paragraph_mut(&mut section.content, para_id) {
                    *section_paragraph = replacement.clone();
                }
            }
        }
        self.authored_ids.retain(|_, (fresh, _)| fresh != para_id);
        Ok(receipt)
    }

    /// Exports the current model as read-only structured content. Anchors address the returned
    /// snapshot; source-part provenance is kept only where the opened package still matches.
    pub fn export_structured(&self, options: &ExportOptions) -> Result<DocxStructuredContent> {
        Ok(export_package_structured(
            self.envelope(),
            &self.original_parts,
            options,
        )?)
    }

    /// [`Document::export_structured`] rendered as Markdown.
    pub fn export_markdown(&self, options: &ExportOptions) -> Result<MarkdownContent> {
        let content = self.export_structured(options)?;
        render_docx_markdown(
            &content,
            &MarkdownOptions {
                max_bytes: options.max_bytes,
            },
        )
    }

    /// Lists the content controls of the current model in document order. Control ids and
    /// anchors address the returned snapshot; controls whose source content no longer matches
    /// the opened package cannot be shown fillable. An editing session fills them.
    pub fn list_content_controls(
        &self,
        options: &ContentControlsOptions,
    ) -> Result<ContentControlsSnapshot> {
        Ok(list_package_content_controls(
            self.envelope(),
            &self.original_parts,
            options,
        )?)
    }

    /// The content controls of the current model that match `query` exactly.
    pub fn find_content_controls(
        &self,
        query: &ContentControlQuery,
        options: &ContentControlsOptions,
    ) -> Result<ContentControlsSnapshot> {
        Ok(find_package_content_controls(
            self.envelope(),
            &self.original_parts,
            query,
            options,
        )?)
    }

    /// The current model as the parsed package seeding reads.
    fn envelope(&self) -> S9WireEnvelope {
        let model = self.model.clone();
        S9WireEnvelope {
            wire_version: 1,
            document: S9DocumentWire {
                package: S9PackageWire {
                    document: S9DocumentBodyWire::from(self.with_authored_ids(model.body)),
                    styles: model.styles,
                    theme: model.theme,
                    numbering: model.numbering,
                    settings: model.settings,
                    font_table: model.font_table,
                    header_entries: Some(model.headers),
                    footer_entries: Some(model.footers),
                    footnotes: Some(model.footnotes),
                    endnotes: Some(model.endnotes),
                    footnote_separators: Some(model.footnote_separators),
                    endnote_separators: Some(model.endnote_separators),
                    relationship_entries: model.relationships,
                    media_entries: Vec::new(),
                    chart_entries: model.charts,
                },
                template_variables: None,
                warnings: (!model.warnings.is_empty()).then_some(model.warnings),
            },
            embedded_font_parts: Vec::new(),
            font_table_relationships_xml: None,
            canonical_base64: None,
            canonical_sha256: None,
        }
    }

    pub fn layout(&self, mut input: LayoutInput) -> Result<LayoutResult> {
        let layout = docx_layout::compute_layout_input(&mut input)?;
        let display_list =
            docx_layout::build_display_list(&input, &layout).map_err(Error::DisplayList)?;
        Ok(LayoutResult {
            layout,
            display_list,
        })
    }

    pub fn save(&self) -> Result<Vec<u8>> {
        self.save_with_options(SaveOptions::default())
    }

    pub fn save_with_options(&self, options: SaveOptions) -> Result<Vec<u8>> {
        let request = S13SaveRequest {
            determinism: SerializerDeterminism {
                seed: self.seed.clone(),
                now: options.now,
            },
            document: self.with_authored_ids(self.model.body.clone()),
            header_entries: self.model.headers.clone(),
            footer_entries: self.model.footers.clone(),
            footnotes: self.model.footnotes.clone(),
            endnotes: self.model.endnotes.clone(),
            footnote_separators: self.model.footnote_separators.clone(),
            endnote_separators: self.model.endnote_separators.clone(),
            relationship_entries: self.model.relationships.clone(),
            numbering: Some(self.model.numbering.clone()),
            options: S13SaveOptions {
                update_modified_date: options.update_modified_date,
                modified_by: options.modified_by,
            },
            selective: None,
            paragraph_ids: None,
        };
        write_docx_s13_parts(request, &self.original_parts, None).map_err(Error::from)
    }

    /// `body` as it saves: a paragraph still at its ordinal with its fresh ID
    /// keeps its authored ID.
    fn with_authored_ids(&self, mut body: DocumentBody) -> DocumentBody {
        if !self.authored_ids.is_empty() {
            let mut ordinal = 0usize;
            visit_paragraphs_mut(&mut body.content, &mut |paragraph| {
                ordinal += 1;
                let Some((fresh, authored)) = self.authored_ids.get(&(ordinal - 1)) else {
                    return;
                };
                if paragraph.para_id.as_ref() == Some(fresh) {
                    let paragraph = Arc::make_mut(paragraph);
                    paragraph.para_id = Some(authored.clone());
                    paragraph.repeated_para_id = Some(true);
                }
            });
        }
        body
    }
}

/// Gives each paragraph that repeats an earlier paragraph's ID a fresh one no
/// package part uses, returning the fresh and authored IDs by paragraph ordinal.
fn address_repeated_paragraphs(
    blocks: &mut [BlockContent],
    parts: &[(String, Vec<u8>)],
) -> BTreeMap<usize, (String, String)> {
    let mut authored_ids = BTreeMap::new();
    let mut occupied = None;
    let mut next_ordinal = 0usize;
    visit_paragraphs_mut(blocks, &mut |paragraph| {
        let ordinal = next_ordinal;
        next_ordinal += 1;
        if paragraph.repeated_para_id != Some(true) {
            return;
        }
        let occupied = occupied.get_or_insert_with(|| package_paragraph_ids(parts));
        let Some(id) = allocate_paragraph_id(&format!("body#{ordinal}"), occupied) else {
            return;
        };
        occupied.insert(id);
        let paragraph = Arc::make_mut(paragraph);
        let fresh = format_paragraph_id(id);
        if let Some(authored) = paragraph.para_id.replace(fresh.clone()) {
            authored_ids.insert(ordinal, (fresh, authored));
        }
        paragraph.repeated_para_id = None;
    });
    authored_ids
}

fn visit_paragraphs_mut(blocks: &mut [BlockContent], visit: &mut impl FnMut(&mut Arc<Paragraph>)) {
    for block in blocks {
        match block {
            BlockContent::Paragraph(paragraph) => visit(paragraph),
            BlockContent::Table(table) => {
                for row in &mut Arc::make_mut(table).rows {
                    for cell in &mut row.cells {
                        visit_paragraphs_mut(&mut cell.content, visit);
                    }
                }
            }
            BlockContent::BlockSdt(sdt) => {
                visit_paragraphs_mut(&mut Arc::make_mut(sdt).content, visit)
            }
            BlockContent::RawXml(_) => {}
        }
    }
}

/// Renders structured content as Markdown within `options.max_bytes`.
pub fn render_docx_markdown(
    content: &DocxStructuredContent,
    options: &MarkdownOptions,
) -> Result<MarkdownContent> {
    Ok(docx_edit::structured::render_docx_markdown(
        content, options,
    )?)
}

type RunTemplate = Option<(
    Option<docx_parse::TextFormatting>,
    Option<Vec<docx_parse::inline::RunPropertyChange>>,
)>;

fn plain_run_template(paragraph: &Paragraph) -> Option<RunTemplate> {
    if paragraph.content.is_empty() {
        return Some(None);
    }
    let [ParagraphContent::Inline(InlineNode::Run(run))] = paragraph.content.as_slice() else {
        return None;
    };
    if !run
        .content
        .iter()
        .all(|content| matches!(content, RunContent::Text { .. }))
    {
        return None;
    }
    Some(Some((run.formatting.clone(), run.property_changes.clone())))
}

fn replace_plain_run(paragraph: &mut Paragraph, text: String, template: RunTemplate) {
    if text.is_empty() {
        paragraph.content.clear();
        return;
    }
    let (formatting, property_changes) = template.unwrap_or_default();
    paragraph.content = vec![ParagraphContent::Inline(InlineNode::Run(Run {
        node_type: RunType::Run,
        formatting,
        property_changes,
        content: vec![RunContent::Text {
            preserve_space: (text.trim() != text).then_some(true),
            text,
        }],
    }))];
}

fn model_from_package(
    package: S9PackageWire,
    template_variables: Vec<String>,
    warnings: Vec<String>,
) -> DocumentModel {
    let S9PackageWire {
        document,
        styles,
        theme,
        numbering,
        settings,
        font_table,
        header_entries,
        footer_entries,
        footnotes,
        endnotes,
        footnote_separators,
        endnote_separators,
        relationship_entries,
        media_entries,
        chart_entries,
    } = package;
    DocumentModel {
        body: body_from_wire(document),
        styles,
        theme,
        numbering,
        settings,
        font_table,
        headers: header_entries.unwrap_or_default(),
        footers: footer_entries.unwrap_or_default(),
        footnotes: footnotes.unwrap_or_default(),
        endnotes: endnotes.unwrap_or_default(),
        footnote_separators: footnote_separators.unwrap_or_default(),
        endnote_separators: endnote_separators.unwrap_or_default(),
        relationships: relationship_entries,
        media: media_entries
            .into_iter()
            .map(|(key, file)| {
                (
                    key,
                    Arc::try_unwrap(file).unwrap_or_else(|file| (*file).clone()),
                )
            })
            .collect(),
        charts: chart_entries,
        template_variables,
        warnings,
    }
}

fn body_from_wire(wire: S9DocumentBodyWire) -> DocumentBody {
    let sections = wire.sections.map(|sections| {
        sections
            .into_iter()
            .map(|section| section_from_wire(section, &wire.content))
            .collect()
    });
    DocumentBody {
        content: wire.content,
        sections,
        final_section_properties: wire.final_section_properties,
        custom_root_bindings: wire.custom_root_bindings,
        comments: wire.comments,
    }
}

fn section_from_wire(section: S9SectionWire, content: &[BlockContent]) -> Section {
    Section {
        id: section.id,
        properties: section.properties,
        content: content
            .get(section.content_start..section.content_end)
            .unwrap_or_default()
            .to_vec(),
    }
}

fn collect_paragraphs<'a>(blocks: &'a [BlockContent], output: &mut Vec<&'a Paragraph>) {
    for block in blocks {
        match block {
            BlockContent::Paragraph(paragraph) => output.push(paragraph),
            BlockContent::Table(table) => {
                for row in &table.rows {
                    for cell in &row.cells {
                        collect_paragraphs(&cell.content, output);
                    }
                }
            }
            BlockContent::BlockSdt(sdt) => collect_paragraphs(&sdt.content, output),
            BlockContent::RawXml(_) => {}
        }
    }
}

fn collect_tables<'a>(blocks: &'a [BlockContent], output: &mut Vec<&'a Table>) {
    for block in blocks {
        match block {
            BlockContent::Paragraph(_) => {}
            BlockContent::Table(table) => {
                output.push(table);
                for row in &table.rows {
                    for cell in &row.cells {
                        collect_tables(&cell.content, output);
                    }
                }
            }
            BlockContent::BlockSdt(sdt) => collect_tables(&sdt.content, output),
            BlockContent::RawXml(_) => {}
        }
    }
}

fn find_paragraph_mut<'a>(
    blocks: &'a mut [BlockContent],
    para_id: &str,
) -> Option<&'a mut Paragraph> {
    for block in blocks {
        match block {
            BlockContent::Paragraph(paragraph) if paragraph.para_id.as_deref() == Some(para_id) => {
                return Some(Arc::make_mut(paragraph));
            }
            BlockContent::Paragraph(_) => {}
            BlockContent::Table(table) => {
                for row in &mut Arc::make_mut(table).rows {
                    for cell in &mut row.cells {
                        if let Some(paragraph) = find_paragraph_mut(&mut cell.content, para_id) {
                            return Some(paragraph);
                        }
                    }
                }
            }
            BlockContent::BlockSdt(sdt) => {
                if let Some(paragraph) =
                    find_paragraph_mut(&mut Arc::make_mut(sdt).content, para_id)
                {
                    return Some(paragraph);
                }
            }
            BlockContent::RawXml(_) => {}
        }
    }
    None
}

fn utf16_len(text: &str) -> u32 {
    text.encode_utf16().count() as u32
}