# shout out to FriendlyWithMeat

from PIL import Image
from pathlib import Path
import os
from PIL.ImageOps import invert
import argparse

parser = argparse.ArgumentParser(
                    prog='invert.py',
                    description='Took images from INPUT_FOLDER, invert them and save in OUTPUT_FOLDER.')
parser.add_argument('INPUT_FOLDER')
parser.add_argument('OUTPUT_FOLDER')

args = parser.parse_args()

WORK_DIR = Path(args.INPUT_FOLDER)
IMG_EXTS = set((".png",))
OUT_DIR = Path(args.OUTPUT_FOLDER)


OUT_DIR.mkdir(exist_ok=True)
dir_contents = os.listdir(WORK_DIR)
for elem in dir_contents:
    img_file = Path(os.path.join(WORK_DIR, elem))
    if img_file.suffix in IMG_EXTS:
        image = Image.open(img_file)
        r, g, b, a = image.convert('RGBA').split()
        rgb_image = Image.merge('RGB', (r,g,b))
        inv_image = invert(rgb_image)
        r, g, b = inv_image.split()
        final_img = Image.merge('RGBA', (r,g,b,a))
        save_path = Path(os.path.join(OUT_DIR, elem))
        final_img.save(save_path)