#!/usr/bin/env python3

import argparse
from pypdf import PaperSize, PdfReader, PdfWriter, Transformation
from pypdf.annotations import Line
from pypdf.generic import ArrayObject, FloatObject, NameObject
import math


def append_folio_A6(input_pdf, offset, size, append_to):
    # Handles edge case where folio has 1 page
    num_A5 = min(math.ceil(size / 4.0) * 2, size)
    writer = append_to
    w_offset = len(writer.pages)
    for i in range(num_A5):
        dpage = writer.add_blank_page(width=PaperSize.A5.height,
                                      height=PaperSize.A5.width)
        ipage = input_pdf.pages[i+offset]
        ipage.mediabox = dpage.mediabox

        if i % 2 == 0: # Even pages are filled right to left
            ipage.add_transformation(Transformation().translate(PaperSize.A6.width,
                                                                0))
        dpage.merge_page(ipage)

    # Now for the fun part: work backwards to fill the remaining pages
    for j in range(size - num_A5):
        dpage = writer.pages[w_offset + num_A5 - 1 - j]
        ipage = input_pdf.pages[num_A5+j+offset]
        ipage.mediabox = dpage.mediabox
        if (num_A5 - 1 - j) % 2 == 1:
            ipage.add_transformation(Transformation().translate(PaperSize.A6.width,
                                                                0))
        dpage.merge_page(ipage)
    return writer


# Interprets input_pdf as a list of A6 pages, and appends to append_to with
# A5 pages using the following rules:
#   a. Odd/even page pairs correspond to front/back pages of a single sheet, printed
#      with "mirror around **short** edge"
#   b. when grouped as consecutive sequences of max_folio_size or less, the printed
#      sheets can be folded in the middle and assembled as a folio
def foliate_A6(input_pdf, max_folio_size=-1, append_to=PdfWriter()):
    current_page = 0
    doc_length = len(input_pdf.pages)
    writer = append_to
    if max_folio_size < 0:
        max_folio_size = doc_length

    folio_page_size = 0
    while(current_page < doc_length):
        # max_folio_size is a number of _sheets_, each containing 4 pages
        folio_page_size = min(doc_length - current_page, 4*max_folio_size)
        append_folio_A6(input_pdf, current_page, folio_page_size, append_to)
        current_page += folio_page_size

    print("Dangling folio has",folio_page_size,"pages",
          "(",int(math.ceil(folio_page_size/4.0)),"sheets).")

    return writer

# Take an A5 printing plan, with odd/even pages being understood
# as front/back pairs of sheets, short edge mirroring. Produces an A4 printing plan
# containing two A5 panels per sheet, long edge mirroring.
def blit_to_A4(input_pdf, append_to=PdfWriter()):
    L = len(input_pdf.pages)
    writer = append_to

    for k in range(0, L, 2):
        if(k%4 == 0): # Open up a new sheet
            front_page = writer.add_blank_page(width=PaperSize.A4.width,
                                               height=PaperSize.A4.height)
            back_page = writer.add_blank_page(width=PaperSize.A4.width,
                                              height=PaperSize.A4.height)

            midline = Line(rect = (0, PaperSize.A4.height/2, PaperSize.A4.width, PaperSize.A4.height/2),
                           p1=(10, PaperSize.A4.height/2),
                           p2=(PaperSize.A4.width-10, PaperSize.A4.height/2))
            midline.flags = 4 # print the midline
            midline[NameObject("/C")] = ArrayObject([FloatObject(0.5)])
            writer.add_annotation(front_page, midline)

        front = input_pdf.pages[k]
        back = input_pdf.pages[k+1] if (k+1<L) else create_blank_page(None, PaperSize.A5.width, PaperSize.A5.height)

        front.mediabox = front_page.mediabox
        back.mediabox  = back_page.mediabox

        if(k%4 == 0): # Translate to bottom half of page
            front.add_transformation(Transformation().translate(0, PaperSize.A5.width))
            back.add_transformation(Transformation().translate(0, PaperSize.A5.width))

        front_page.merge_page(front)
        back_page.merge_page(back)
    return writer

parser = argparse.ArgumentParser()
parser.add_argument("infile",  help="input document, A6 pages")
parser.add_argument("outfile", help="output document")
parser.add_argument("-a5", help="output as a5, reverse on short edge",
                    action="store_true")
parser.add_argument("-fs", help="maximum folio size, in sheets",
                    action="store", default=-1, type=int)
args = parser.parse_args()

input_pdf = PdfReader(args.infile)
a5 = foliate_A6(input_pdf, args.fs)
if (args.a5):
    a5.write(args.outfile)
else:
    a4 = blit_to_A4(a5)
    a4.write(args.outfile)