# image dithering script
# © 2022 Roel Roscam Abbing, released as AGPLv3
# see https://www.gnu.org/licenses/agpl-3.0.html
# Support your local low-tech magazine: https://solar.lowtechmagazine.com/donate.html

import argparse
import logging
import os
import shutil

import hitherdither
from PIL import Image

COLOR = {
    "blue": hitherdither.palette.Palette(
        [
            (30, 32, 40),
            (11, 21, 71),
            (57, 77, 174),
            (158, 168, 218),
            (187, 196, 230),
            (243, 244, 250),
        ]
    ),
    "green": hitherdither.palette.Palette(
        [
            (9, 74, 58),
            (58, 136, 118),
            (101, 163, 148),
            (144, 189, 179),
            (169, 204, 195),
            (242, 247, 246),
        ]
    ),
    "red": hitherdither.palette.Palette(
        [
            (86, 9, 6),
            (197, 49, 45),
            (228, 130, 124),
            (233, 155, 151),
            (242, 193, 190),
            (252, 241, 240),
        ]
    ),
    "grayscale": hitherdither.palette.Palette(
        [
            (25, 25, 25),
            (75, 75, 75),
            (125, 125, 125),
            (175, 175, 175),
            (225, 225, 225),
            (250, 250, 250),
        ]
    ),
}

parser = argparse.ArgumentParser(
    """
    This script recursively traverses folders and creates dithered versions of the images it finds.
    These are stored in the same folder as the images in a folder called "dithers".
    """
)

parser.add_argument(
    "-d", "--directory", help="Set the directory to traverse", default="."
)

parser.add_argument(
    "-o", "--output", help="Set the directory to output", default="dithered"
)

parser.add_argument(
    "-rm",
    "--remove",
    help="Removes all the folders with dithers and their contents",
    action="store_true",
)

parser.add_argument(
    "-c", "--colorize", help="Colorizes the dithered images", default="grayscale"
)

parser.add_argument(
    "-v",
    "--verbose",
    help="Print out more detailed information about what this script is doing",
    action="store_true",
)

args = parser.parse_args()

image_ext = [".jpg", ".JPG", ".jpeg", ".png", ".gif", ".webp", ".tiff", ".bmp"]


content_dir = args.directory
output_dir = args.output
os.makedirs(args.output, exist_ok=True)

if args.verbose:
    logging.basicConfig(level=logging.DEBUG)
else:
    logging.basicConfig(level=logging.INFO)

logging.info(
    f"Dithering all images in {content_dir} and subfolders, excluding {output_dir}."
)


def dither_image(source_image, output_image):
    # see hitherdither docs for different dithering algos and settings

    palette = COLOR[args.colorize]
    try:
        img = Image.open(source_image).convert("RGB")
        img.thumbnail((800, 800), Image.BICUBIC)
        threshold = [96, 96, 96]
        img_dithered = hitherdither.ordered.bayer.bayer_dithering(
            img, palette, threshold, order=4
        )  # 8
        img_dithered.save(output_image, optimize=True, compress_level=9, quality=0)

    except Exception as e:
        logging.debug("❌ failed to convert {}".format(source_image))
        logging.debug(e)


def delete_dithers(content_dir, output_dir):
    logging.info(f"Deleting '{output_dir}' folders in {content_dir} and below")
    for root, dirs, files in os.walk(content_dir, topdown=True):
        if root.endswith(output_dir):
            shutil.rmtree(root)
            logging.info("Removed {}".format(root))


prev_root = None

if args.remove:
    delete_dithers(
        os.path.abspath(content_dir),
        os.path.abspath(output_dir),
    )
else:
    for root, dirs, files in os.walk(os.path.abspath(content_dir), topdown=True):
        logging.debug("Checking next folder {}".format(root))

        print(dirs, output_dir)
        dirs[:] = [d for d in dirs if d != output_dir]

        if prev_root is None:
            prev_root = root

        if prev_root is not root:
            if files:
                if any(x.endswith(tuple(image_ext)) for x in files):
                    if not os.path.exists(os.path.join(root, output_dir)):
                        os.mkdir(os.path.join(root, output_dir))
                        logging.info("📁 created in {}".format(root))

        for fname in files:
            if fname.endswith(tuple(image_ext)):
                file_, ext = os.path.splitext(fname)
                source_image = os.path.join(root, fname)
                output_image = os.path.join(
                    os.path.join(root, output_dir), file_ + "_dithered.png"
                )
                if not os.path.exists(output_image):
                    dither_image(source_image, output_image)
                    logging.info("🖼 converted {}".format(fname))
                    logging.debug(output_image)
                else:
                    logging.debug(
                        "Dithered version of {} found, skipping".format(fname)
                    )

        prev_root = root

logging.info("Done dithering")