merge_pdf

use std::{collections::BTreeMap, io::Cursor};

use lopdf::{Document, Object, ObjectId};

use rustler::{Binary, NewBinary, OwnedBinary};

enum Error {
    Lopdf(lopdf::Error),
    Merge(String),
}

impl From<lopdf::Error> for Error {
    fn from(err: lopdf::Error) -> Error {
        Error::Lopdf(err)
    }
}

impl rustler::Encoder for Error {
    fn encode<'a>(&self, env: rustler::Env<'a>) -> rustler::Term<'a> {
        let msg = match self {
            Error::Lopdf(_err) => "lopdf error",
            Error::Merge(err) => err,
        };

        let mut msg_binary = NewBinary::new(env, msg.len());
        msg_binary.as_mut_slice().copy_from_slice(msg.as_bytes());

        msg_binary.into()
    }
}

#[rustler::nif(schedule = "DirtyCpu")]
fn merge_paths(env: rustler::Env, paths: Vec<String>) -> Result<Binary, Error> {
    let mut documents: Vec<Document> = vec![];
    for path in paths {
        let document = Document::load(&path)?;
        documents.push(document);
    }

    let output_vec = merge_documents(documents)?;
    let output_value = output_vec.get_ref();

    let mut output_binary = OwnedBinary::new(output_value.len()).unwrap();
    output_binary.as_mut_slice().copy_from_slice(output_value);

    Ok(Binary::from_owned(output_binary, env))
}

#[rustler::nif(schedule = "DirtyCpu")]
fn merge_binaries<'a>(
    env: rustler::Env<'a>,
    binaries: Vec<Binary<'a>>,
) -> Result<Binary<'a>, Error> {
    let mut documents: Vec<Document> = vec![];
    for binary in binaries {
        let document = Document::load_mem(&binary.as_slice())?;
        documents.push(document);
    }

    let output_vec = merge_documents(documents)?;
    let output_value = output_vec.get_ref();

    let mut output_binary = OwnedBinary::new(output_value.len()).unwrap();
    output_binary.as_mut_slice().copy_from_slice(output_value);

    Ok(Binary::from_owned(output_binary, env))
}

fn merge_documents(documents: Vec<Document>) -> Result<Cursor<Vec<u8>>, Error> {
    // Define a starting max_id (will be used as start index for object_ids)
    let mut max_id = 1;
    let mut pagenum = 1;
    // Collect all Documents Objects grouped by a map
    let mut documents_pages = BTreeMap::new();
    let mut documents_objects = BTreeMap::new();
    let mut document = Document::with_version("1.5");

    for mut doc in documents {
        let mut first = false;
        doc.renumber_objects_with(max_id);

        max_id = doc.max_id + 1;

        documents_pages.extend(
            doc.get_pages()
                .into_iter()
                .map(|(_, object_id)| {
                    if !first {
                        // let bookmark = Bookmark::new(
                        //     String::from(format!("Page_{}", pagenum)),
                        //     [0.0, 0.0, 1.0],
                        //     0,
                        //     object_id,
                        // );
                        // document.add_bookmark(bookmark, None);
                        first = true;
                        pagenum += 1;
                    }
                    (object_id, doc.get_object(object_id).unwrap().to_owned())
                })
                .collect::<BTreeMap<ObjectId, Object>>(),
        );
        documents_objects.extend(doc.objects);
    }

    // Catalog and Pages are mandatory
    let mut catalog_object: Option<(ObjectId, Object)> = None;
    let mut pages_object: Option<(ObjectId, Object)> = None;

    // Process all objects except "Page" type
    for (object_id, object) in documents_objects.iter() {
        // We have to ignore "Page" (as are processed later), "Outlines" and "Outline" objects
        // All other objects should be collected and inserted into the main Document
        match object.type_name().unwrap_or("") {
            "Catalog" => {
                // Collect a first "Catalog" object and use it for the future "Pages"
                catalog_object = Some((
                    if let Some((id, _)) = catalog_object {
                        id
                    } else {
                        *object_id
                    },
                    object.clone(),
                ));
            }
            "Pages" => {
                // Collect and update a first "Pages" object and use it for the future "Catalog"
                // We have also to merge all dictionaries of the old and the new "Pages" object
                if let Ok(dictionary) = object.as_dict() {
                    let mut dictionary = dictionary.clone();
                    if let Some((_, ref object)) = pages_object {
                        if let Ok(old_dictionary) = object.as_dict() {
                            dictionary.extend(old_dictionary);
                        }
                    }

                    pages_object = Some((
                        if let Some((id, _)) = pages_object {
                            id
                        } else {
                            *object_id
                        },
                        Object::Dictionary(dictionary),
                    ));
                }
            }
            "Page" => {}     // Ignored, processed later and separately
            "Outlines" => {} // Ignored, not supported yet
            "Outline" => {}  // Ignored, not supported yet
            _ => {
                document.objects.insert(*object_id, object.clone());
            }
        }
    }

    // If no "Pages" object found abort
    if pages_object.is_none() {
        return Err(Error::Merge("Pages root not found.".to_string()));
    }

    // Iterate over all "Page" objects and collect into the parent "Pages" created before
    for (object_id, object) in documents_pages.iter() {
        if let Ok(dictionary) = object.as_dict() {
            let mut dictionary = dictionary.clone();
            dictionary.set("Parent", pages_object.as_ref().unwrap().0);

            document
                .objects
                .insert(*object_id, Object::Dictionary(dictionary));
        }
    }

    // If no "Catalog" found abort
    if catalog_object.is_none() {
        return Err(Error::Merge("No catalog".to_string()));
    }

    let catalog_object = catalog_object.unwrap();
    let pages_object = pages_object.unwrap();

    // Build a new "Pages" with updated fields
    if let Ok(dictionary) = pages_object.1.as_dict() {
        let mut dictionary = dictionary.clone();

        // Set new pages count
        dictionary.set("Count", documents_pages.len() as u32);

        // Set new "Kids" list (collected from documents pages) for "Pages"
        dictionary.set(
            "Kids",
            documents_pages
                .into_iter()
                .map(|(object_id, _)| Object::Reference(object_id))
                .collect::<Vec<_>>(),
        );

        document
            .objects
            .insert(pages_object.0, Object::Dictionary(dictionary));
    }

    // Build a new "Catalog" with updated fields
    if let Ok(dictionary) = catalog_object.1.as_dict() {
        let mut dictionary = dictionary.clone();
        dictionary.set("Pages", pages_object.0);
        dictionary.remove(b"Outlines"); // Outlines not supported in merged PDFs

        document
            .objects
            .insert(catalog_object.0, Object::Dictionary(dictionary));
    }

    document.trailer.set("Root", catalog_object.0);

    // Update the max internal ID as wasn't updated before due to direct objects insertion
    document.max_id = document.objects.len() as u32;

    // Reorder all new Document objects
    document.renumber_objects();

    //Set any Bookmarks to the First child if they are not set to a page
    document.adjust_zero_pages();

    //Set all bookmarks to the PDF Object tree then set the Outlines to the Bookmark content map.
    if let Some(n) = document.build_outline() {
        if let Ok(x) = document.get_object_mut(catalog_object.0) {
            if let Object::Dictionary(ref mut dict) = x {
                dict.set("Outlines", Object::Reference(n));
            }
        }
    }

    document.compress();

    // Save the merged PDF
    // Store file in current working directory.
    // Note: Line is excluded when running tests
    let mut buff: Cursor<Vec<u8>> = Cursor::new(Vec::new());
    match document.save_to(&mut buff) {
        Ok(_) => {
            return Ok(buff);
        }
        Err(_) => return Err(Error::Merge("Failed to save".to_string())),
    }
}

rustler::init!("Elixir.MergePdf.Native", [merge_binaries, merge_paths]);