#!/usr/bin/env cached-nix-shell
#!nix-shell -i python3 -p python3 python3Packages.pillow python3Packages.tqdm
import argparse
import sys
import zlib
from concurrent.futures import ThreadPoolExecutor
from zipfile import ZipFile

from PIL import Image, ImageOps
from tqdm import tqdm


def process_image(im):
    # double pages (wider than high) should be rotated
    if im.width > im.height:
        im = im.transpose(Image.Transpose.ROTATE_90)

    # convert to greyscale
    im = im.convert("L")

    # pad to screen size
    im = ImageOps.pad(im, (args.width, args.height), Image.Resampling.LANCZOS, "#fff")

    return im


def process_archive(srczip, dstzip, multithread=True):
    executor = ThreadPoolExecutor()

    if multithread:
        mapper = executor.map
    else:
        mapper = map

    srcfiles = sorted(
        map(
            lambda zi: zi.filename,
            filter(lambda zi: not zi.is_dir(), srczip.infolist()),
        )
    )

    for idx, im in enumerate(
        tqdm(
            mapper(process_image, mapper(Image.open, mapper(srczip.open, srcfiles))),
            total=len(srcfiles),
        )
    ):
        with dstzip.open(f"{idx+1:04d}.jpg", "w") as outpagefile:
            im.save(outpagefile, format="JPEG", quality=92)


parser = argparse.ArgumentParser()
parser.add_argument("infile")
parser.add_argument("outfile")
parser.add_argument("-W", "--width", default=1440)
parser.add_argument("-H", "--height", default=1920)
parser.add_argument("-s", "--singlethreaded", action="store_true")
args = parser.parse_args()

srczip = ZipFile(args.infile)
dstzip = ZipFile(args.outfile, "x")

if args.singlethreaded:
    process_archive(srczip, dstzip, multithread=False)
else:
    try:
        process_archive(srczip, dstzip, multithread=True)
    except zlib.error as e:
        print(f"zlib choked ({e}), retrying in single threaded mode", file=sys.stderr)
        # replace partial archive with new empty one
        dstzip.close()
        dstzip = ZipFile(args.outfile, "w")
        process_archive(srczip, dstzip, multithread=False)