#!/usr/bin/env python3
import argparse

from pypdf import PdfReader, PdfWriter, Transformation, PageObject, PaperSize


def double_folio(source, out, section_length, outsize="A4", extra_scale=1):
    assert section_length % 4 == 0
    in_reader = PdfReader(source)
    out_writer = PdfWriter()

    to_add = section_length - len(in_reader.pages) % section_length
    add_before = to_add // 2
    add_after = to_add - add_before

    out_size = getattr(PaperSize, outsize)

    pages = [PageObject.create_blank_page(width=out_size.width, height=out_size.height) for _ in range(add_before)]
    pages += in_reader.pages
    pages += [PageObject.create_blank_page(width=out_size.width, height=out_size.height) for _ in range(add_after)]

    for sig_start in range(0, len(in_reader.pages), section_length):
        for side in range(section_length // 4):
            top = [
                sig_start + section_length - side - 1,
                sig_start + side,
            ]
            bottom = [
                sig_start + (section_length // 2) + side,
                sig_start + (section_length // 2) - side - 1
            ]

            if side % 2 != 0:
                top = list(reversed(top))
                bottom = list(reversed(bottom))

            imposed_page = PageObject.create_blank_page(width=out_size.width, height=out_size.height)

            scale = Transformation().scale(0.5 * extra_scale, 0.5 * extra_scale)
            up = Transformation().translate(ty=extra_scale * out_size.height / 2)
            right = Transformation().translate(tx=extra_scale * out_size.width / 2)
            rot = Transformation().rotate(180)

            tl = pages[top[0]]
            tl.add_transformation(scale)
            tl.add_transformation(up)
            imposed_page.merge_page(tl)

            tr = pages[top[1]]
            tr.add_transformation(scale)
            tr.add_transformation(up)
            tr.add_transformation(right)
            imposed_page.merge_page(tr)

            bl = pages[bottom[0]]
            bl.add_transformation(scale)
            bl.add_transformation(rot)
            bl.add_transformation(up)
            bl.add_transformation(right)
            imposed_page.merge_page(bl)

            br = pages[bottom[1]]
            br.add_transformation(scale)
            br.add_transformation(rot)
            br.add_transformation(right)
            br.add_transformation(right)
            br.add_transformation(up)
            imposed_page.merge_page(br)

            out_writer.add_page(imposed_page)

    with open(out, "wb") as fp:
        out_writer.write(fp)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="pympose",
        description="Impose pages for printing"
    )

    parser.add_argument('action', type=str, help="What to do")
    parser.add_argument('infile', type=str, help="Input file")
    parser.add_argument('outfile', type=str, help="Output file")
    parser.add_argument('section_size', type=int, help="Size of sections")
    parser.add_argument('--outsize', type=str, help="Output Size", default="A4")
    parser.add_argument('--extra-scale', type=float, help="Extra Scale factor", default=1)

    args = parser.parse_args()
    if args.action == "double_folio":
        double_folio(args.infile, args.outfile, args.section_size, args.outsize, args.extra_scale)