djvu-rs 0.20.0

//! DJVM document merge and split operations.
//!
//! Provides [`merge`] to combine multiple DjVu documents into a single
//! bundled DJVM, and [`split`] to extract page ranges from a document.
//!
//! [`merge`]: crate::djvm::merge
//! [`split`]: crate::djvm::split

#[cfg(not(feature = "std"))]
use alloc::{format, string::String, vec, vec::Vec};

use crate::dirm::DirmPayload;
use crate::error::IffError;
use crate::iff;

#[cfg(test)]
use crate::djvu_document::DjVuDocument;

/// Error type for merge/split operations.
#[derive(Debug, thiserror::Error)]
pub enum DjvmError {
    /// IFF container parse error.
    #[error("IFF parse error: {0}")]
    Iff(#[from] IffError),

    /// Document model error.
    #[error("document error: {0}")]
    Doc(#[from] crate::djvu_document::DocError),

    /// No pages to merge.
    #[error("no pages to merge")]
    EmptyMerge,

    /// Page range is out of bounds.
    #[error("page range {start}..{end} is out of bounds (document has {count} pages)")]
    PageRangeOutOfBounds {
        start: usize,
        end: usize,
        count: usize,
    },
}

/// Re-serialize a sub-FORM child — the raw `data` of a `FORM` chunk, which
/// begins with its 4-byte form type — back into a standalone `AT&T`-prefixed
/// FORM document. Inverse of [`strip_att`].
fn wrap_sub_form(form_data: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(12 + form_data.len());
    out.extend_from_slice(b"AT&T");
    out.extend_from_slice(b"FORM");
    out.extend_from_slice(&(form_data.len() as u32).to_be_bytes());
    out.extend_from_slice(form_data);
    out
}

/// Strip a leading `AT&T` magic from a standalone FORM document, yielding the
/// `FORM`-chunk bytes to embed inside a DJVM bundle. Inverse of [`wrap_sub_form`].
fn strip_att(form: &[u8]) -> &[u8] {
    if form.len() >= 4 && &form[..4] == b"AT&T" {
        &form[4..]
    } else {
        form
    }
}

/// Merge multiple DjVu documents (raw bytes) into a single bundled DJVM.
///
/// Each input document contributes all its pages to the output.
/// Shared dictionaries (DJVI components) are included and INCL
/// references are preserved within each source document's pages.
pub fn merge(documents: &[&[u8]]) -> Result<Vec<u8>, DjvmError> {
    if documents.is_empty() {
        return Err(DjvmError::EmptyMerge);
    }

    let mut components: Vec<Vec<u8>> = Vec::new();
    let mut component_ids: Vec<String> = Vec::new();
    let mut component_flags: Vec<u8> = Vec::new();

    for (doc_idx, &doc_data) in documents.iter().enumerate() {
        let form = iff::parse_form(doc_data)?;

        if &form.form_type == b"DJVU" {
            // Single-page document — the whole file is one page
            components.push(doc_data.to_vec());
            component_ids.push(format!("p{:04}.djvu", components.len()));
            component_flags.push(1); // page
        } else if &form.form_type == b"DJVM" {
            // Multi-page bundled document — extract each FORM child
            for chunk in &form.chunks {
                if &chunk.id == b"FORM" && chunk.data.len() >= 4 {
                    let child_form_type = &chunk.data[..4];

                    let flag = if child_form_type == b"DJVI" { 0 } else { 1 }; // 0 = shared, 1 = page

                    components.push(wrap_sub_form(chunk.data));
                    component_ids.push(format!("d{}p{:04}.djvu", doc_idx, components.len()));
                    component_flags.push(flag);
                }
            }
        }
    }

    if components.is_empty() {
        return Err(DjvmError::EmptyMerge);
    }

    build_djvm(&components, &component_ids, &component_flags)
}

/// Split a document, extracting pages in the given range (0-based, exclusive end).
///
/// Returns raw DjVu bytes for a new document containing only the requested pages.
pub fn split(doc_data: &[u8], start: usize, end: usize) -> Result<Vec<u8>, DjvmError> {
    let form = iff::parse_form(doc_data)?;

    // Page count derived from the same FORM walk used for extraction below, so
    // the bounds check can never disagree with what is actually present (a
    // DIRM-based page count and the FORM:DJVU children can diverge).
    let count = match &form.form_type {
        b"DJVU" => 1,
        b"DJVM" => form
            .chunks
            .iter()
            .filter(|c| &c.id == b"FORM" && c.data.len() >= 4 && &c.data[..4] == b"DJVU")
            .count(),
        _ => 0,
    };

    if start >= count || end > count || start >= end {
        return Err(DjvmError::PageRangeOutOfBounds { start, end, count });
    }

    // Single-page document: just return the whole thing
    if &form.form_type == b"DJVU" && start == 0 && end == 1 {
        return Ok(doc_data.to_vec());
    }

    // For a single page extraction from a multi-page document
    if end - start == 1 && &form.form_type == b"DJVM" {
        let mut page_idx = 0;
        for chunk in &form.chunks {
            if &chunk.id == b"FORM" && chunk.data.len() >= 4 && &chunk.data[..4] == b"DJVU" {
                if page_idx == start {
                    return Ok(wrap_sub_form(chunk.data));
                }
                page_idx += 1;
            }
        }
    }

    // Multiple pages: build a new DJVM bundle with the requested range
    let mut components: Vec<Vec<u8>> = Vec::new();
    let mut component_ids: Vec<String> = Vec::new();
    let mut component_flags: Vec<u8> = Vec::new();

    // First pass: collect shared components (DJVI) that might be needed
    for chunk in &form.chunks {
        if &chunk.id == b"FORM" && chunk.data.len() >= 4 && &chunk.data[..4] == b"DJVI" {
            components.push(wrap_sub_form(chunk.data));
            component_ids.push(format!("shared{}.djvi", components.len()));
            component_flags.push(0); // shared
        }
    }

    // Second pass: collect pages in the requested range
    let mut page_idx = 0;
    for chunk in &form.chunks {
        if &chunk.id == b"FORM" && chunk.data.len() >= 4 && &chunk.data[..4] == b"DJVU" {
            if page_idx >= start && page_idx < end {
                components.push(wrap_sub_form(chunk.data));
                component_ids.push(format!("p{:04}.djvu", page_idx + 1));
                component_flags.push(1); // page
            }
            page_idx += 1;
        }
    }

    build_djvm(&components, &component_ids, &component_flags)
}

/// Build a bundled DJVM file from components.
fn build_djvm(components: &[Vec<u8>], ids: &[String], flags: &[u8]) -> Result<Vec<u8>, DjvmError> {
    let n = components.len();

    // Build DIRM chunk (bundled; offset slots zeroed — readers fall back to
    // FORM boundaries, matching prior behavior).
    let dirm_data = DirmPayload::build_bundled(n, flags, ids).encode();

    // Calculate total FORM body size
    let mut body_size: usize = 4; // "DJVM"
    body_size += 8 + dirm_data.len(); // DIRM chunk header + data
    if !dirm_data.len().is_multiple_of(2) {
        body_size += 1; // IFF padding
    }
    for comp in components {
        // Each component includes AT&T prefix — strip it for embedding
        let comp_data = strip_att(comp);
        body_size += comp_data.len();
        if !comp_data.len().is_multiple_of(2) {
            body_size += 1; // IFF padding
        }
    }

    let mut output = Vec::with_capacity(4 + 4 + 4 + body_size);

    // AT&T magic
    output.extend_from_slice(b"AT&T");
    // FORM header
    output.extend_from_slice(b"FORM");
    output.extend_from_slice(&(body_size as u32).to_be_bytes());
    // DJVM type
    output.extend_from_slice(b"DJVM");

    // DIRM chunk
    output.extend_from_slice(b"DIRM");
    output.extend_from_slice(&(dirm_data.len() as u32).to_be_bytes());
    output.extend_from_slice(&dirm_data);
    if !dirm_data.len().is_multiple_of(2) {
        output.push(0); // IFF padding
    }

    // Component FORM chunks
    for comp in components {
        let comp_data = strip_att(comp);
        output.extend_from_slice(comp_data);
        if !comp_data.len().is_multiple_of(2) {
            output.push(0); // IFF padding
        }
    }

    Ok(output)
}

/// Create an indirect (non-bundled) DJVM index file that references pages as
/// separate files.
///
/// The returned bytes are a valid `FORM:DJVM` with a DIRM directory chunk whose
/// `is_bundled` flag is **not** set.  Each entry in `page_names` becomes one
/// `Page` component; there are no embedded `FORM:DJVU` sub-forms — the component
/// data lives in separate files that must be passed to a resolver when parsing.
///
/// Shared-dictionary (DJVI) components are not supported by this helper; use
/// [`merge`] to build a bundled document that includes them.
///
/// # Errors
///
/// Returns [`DjvmError::EmptyMerge`] if `page_names` is empty.
pub fn create_indirect(page_names: &[&str]) -> Result<Vec<u8>, DjvmError> {
    if page_names.is_empty() {
        return Err(DjvmError::EmptyMerge);
    }

    let count = page_names.len();
    let ids: Vec<String> = page_names.iter().map(|s| s.to_string()).collect();
    // All entries are pages (flag = 1)
    let flags: Vec<u8> = vec![1u8; count];

    let dirm_data = DirmPayload::build_indirect(count, &flags, &ids).encode();

    let mut body_size: usize = 4; // "DJVM"
    body_size += 8 + dirm_data.len(); // DIRM chunk header + data
    if !dirm_data.len().is_multiple_of(2) {
        body_size += 1;
    }

    let mut output = Vec::with_capacity(4 + 4 + 4 + body_size);
    output.extend_from_slice(b"AT&T");
    output.extend_from_slice(b"FORM");
    output.extend_from_slice(&(body_size as u32).to_be_bytes());
    output.extend_from_slice(b"DJVM");
    output.extend_from_slice(b"DIRM");
    output.extend_from_slice(&(dirm_data.len() as u32).to_be_bytes());
    output.extend_from_slice(&dirm_data);
    if !dirm_data.len().is_multiple_of(2) {
        output.push(0);
    }

    Ok(output)
}

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

    fn fixture_path(name: &str) -> std::path::PathBuf {
        std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("tests/fixtures")
            .join(name)
    }

    #[test]
    fn merge_empty_returns_error() {
        let result = merge(&[]);
        assert!(result.is_err());
    }

    #[test]
    fn split_single_page_from_multipage() {
        let path = fixture_path("DjVu3Spec_bundled.djvu");
        if !path.exists() {
            // Skip if fixture not available
            return;
        }
        let data = std::fs::read(&path).expect("read fixture");
        let doc = DjVuDocument::parse(&data).expect("parse");
        let count = doc.page_count();
        assert!(count > 1, "need multipage fixture");

        // Split out page 0
        let page0 = split(&data, 0, 1).expect("split page 0");
        // Verify the result is parseable
        let form = iff::parse_form(&page0).expect("parse split page");
        assert_eq!(&form.form_type, b"DJVU");
    }

    #[test]
    fn merge_two_single_page_files() {
        let path = fixture_path("irish.djvu");
        if !path.exists() {
            return;
        }
        let irish = std::fs::read(&path).expect("read fixture");
        let data = merge(&[&irish, &irish]).expect("merge");
        // Verify the result has the right FORM type
        let form = iff::parse_form(&data).expect("parse merged");
        assert_eq!(&form.form_type, b"DJVM");
    }

    #[test]
    fn split_out_of_bounds() {
        let path = fixture_path("irish.djvu");
        if !path.exists() {
            return;
        }
        let data = std::fs::read(&path).expect("read fixture");
        let result = split(&data, 0, 5);
        assert!(result.is_err());
    }

    #[test]
    fn create_indirect_empty_returns_error() {
        let result = create_indirect(&[]);
        assert!(result.is_err());
    }

    #[test]
    fn create_indirect_parses_with_resolver() {
        // Build an indirect DJVM that references "chicken.djvu"
        let indirect_bytes = create_indirect(&["chicken.djvu"]).expect("create_indirect");

        // Verify it parses as FORM:DJVM
        let form = iff::parse_form(&indirect_bytes).expect("parse form");
        assert_eq!(&form.form_type, b"DJVM");

        // Verify DIRM chunk has is_bundled = 0
        let dirm = form.chunks.iter().find(|c| &c.id == b"DIRM").expect("DIRM");
        assert_eq!(
            dirm.data[0] & 0x80,
            0,
            "indirect DIRM must not have bundled bit set"
        );

        // Parse with a resolver that supplies chicken.djvu
        let chicken_path = fixture_path("chicken.djvu");
        if !chicken_path.exists() {
            return;
        }
        let chicken_data = std::fs::read(&chicken_path).expect("read chicken.djvu");
        let doc = DjVuDocument::parse_with_resolver(
            &indirect_bytes,
            Some(
                move |name: &str| -> Result<Vec<u8>, crate::djvu_document::DocError> {
                    if name == "chicken.djvu" {
                        Ok(chicken_data.clone())
                    } else {
                        Err(crate::djvu_document::DocError::IndirectResolve(
                            name.to_string(),
                        ))
                    }
                },
            ),
        )
        .expect("parse indirect with resolver");

        assert_eq!(doc.page_count(), 1);
        let page = doc.page(0).unwrap();
        assert_eq!(page.width(), 181);
        assert_eq!(page.height(), 240);
    }

    #[test]
    fn create_indirect_multipage() {
        // 3-page indirect document
        let indirect_bytes =
            create_indirect(&["page1.djvu", "page2.djvu", "page3.djvu"]).expect("create_indirect");
        let form = iff::parse_form(&indirect_bytes).expect("parse");
        assert_eq!(&form.form_type, b"DJVM");

        // Component count = 3 in DIRM
        let dirm = form.chunks.iter().find(|c| &c.id == b"DIRM").expect("DIRM");
        let nfiles = u16::from_be_bytes([dirm.data[1], dirm.data[2]]) as usize;
        assert_eq!(nfiles, 3);
    }

    #[test]
    fn parse_from_dir_indirect() {
        // Write an indirect DJVM index and chicken.djvu to a temp directory,
        // then open it via parse_from_dir.
        let chicken_path = fixture_path("chicken.djvu");
        if !chicken_path.exists() {
            return;
        }
        let tmp = std::env::temp_dir().join("djvu_indirect_test");
        std::fs::create_dir_all(&tmp).unwrap();

        // Copy chicken.djvu as the component
        let component_name = "p0001.djvu";
        std::fs::copy(&chicken_path, tmp.join(component_name)).unwrap();

        // Build indirect index
        let index_bytes = create_indirect(&[component_name]).expect("create_indirect");
        let index_path = tmp.join("index.djvu");
        std::fs::write(&index_path, &index_bytes).unwrap();

        // Open via parse_from_dir
        let index_data = std::fs::read(&index_path).unwrap();
        let doc = DjVuDocument::parse_from_dir(&index_data, &tmp).expect("parse_from_dir");
        assert_eq!(doc.page_count(), 1);
        assert_eq!(doc.page(0).unwrap().width(), 181);
    }
}