#!/usr/bin/env python3

import sys
import os
import json
import time
from datetime import datetime

import PIL.ExifTags
import PIL.Image
from PIL import Image, ImageDraw, ImageFont
import argparse

def yellow(text):
    return "\033[93m" + text + "\033[0m"

def green(text):
    return "\033[92m" + text + "\033[0m"

def red(text):
    return "\033[31m" + text + "\033[0m"

def fetch(att, dict):
    if not att in dict:
        return None
    if dict[att] == None or dict[att] == "":
        return None
    return dict[att]

def pick_in_list(index, arr):
    if (not isinstance(arr, list)) and (not isinstance(arr, tuple)):
        return arr
    if len(arr) > 0:
        return arr[index]
    return None

def pick_ratio(k):
    # 5"      : 127mm * 89mm # NOTE, the real size is 88mm
    # 5" extra: 127mm * 95mm
    # 6"      : 152mm * 102mm
    return {
            "5": 127.0 / 88,
            "5e": 127.0 / 95,
            "6": 152.0 / 102,
            "x": 182.0 / 102
            }[k]

def load_info(img, path, verbose):
    exif = {}
    if (img == None) or\
        (img._getexif() == None) or\
        (img._getexif().items() == None):
        return None
    for k, v in img._getexif().items():
        if (k in PIL.ExifTags.TAGS) and (not isinstance(v, bytes)):
            kk = PIL.ExifTags.TAGS[k]
            exif[kk] = v
    if verbose:
        print(json.dumps(exif, indent = 4))

    info = {}
    # orientation
    ori = fetch("Orientation", exif)
    if ori == 8:
        if verbose:
            print("    rotate img 90'")
        img = img.rotate(90, expand=1)
    # ref https://www.awaresystems.be/imaging/tiff/tifftags/privateifd/exif.html
    # time
    temp = fetch("DateTimeOriginal", exif)
    if temp:
        ct = time.strptime(temp, "%Y:%m:%d %H:%M:%S")
        ct = time.strftime("%Y/%m/%d %H:%M", ct)
    else:
        mtime = os.path.getmtime(path)
        ct = datetime.utcfromtimestamp(mtime).strftime("%Y/%m/%d %H:%M")
    # focal, a list
    temp = fetch("FocalLength", exif)
    if temp:
        focal = int(float(temp[0]) / temp[1])
    else:
        focal = None
    # focal in 35mm
    focal_in35mm = fetch("FocalLengthIn35mmFilm", exif)
    # make
    make = fetch("Make", exif)
    # Model
    model = fetch("Model", exif)
    # ExposureTime
    temp = fetch("ExposureTime", exif)
    if temp:
        value = float(temp[0]) / temp[1]
        if value > 0.4:
            exp_time = "{}".format(value)
        else:
            exp_time = "1/{}".format(int(1.0 / value))
    else:
        exp_time = None
    # f number
    temp = fetch("FNumber", exif)
    if temp:
        temp = float(temp[0]) / temp[1]
        if temp > 8:
            fnumber = int(temp)
        else:
            fnumber = "{0:.1f}".format(temp)
    else:
        fnumber = None
    # WhiteBalance
    """
    Indicates the white balance mode set when the image was shot.
    0 = Auto white balance
    1 = Manual white balance
    """
    temp = fetch("WhiteBalance", exif)
    # exposure bias
    """
    The exposure bias.
    The unit is the APEX (Additive System of Photographic Exposure) setting.
    Ordinarily it is given in the range of -99.99 to 99.99.
    """
    temp = fetch("ExposureBiasValue", exif)
    if temp:
        if temp[0] == 0:
            ev = ""
        else:
            temp = float(temp[0]) / temp[1]
            ev = ", "
            if temp > 0:
                ev += "+"
            if temp < 0:
                ev += "-"
                temp = abs(temp)
            ev += "{0:.1f}ev".format(temp)
    else:
        ev = None
    # exp program
    """
    The class of the program used by the camera to set exposure when the picture is taken.
    0 = Not defined
    1 = Manual
    2 = Normal program
    3 = Aperture priority
    4 = Shutter priority
    5 = Creative program (biased toward depth of field)
    6 = Action program (biased toward fast shutter speed)
    7 = Portrait mode (for closeup photos with the background out of focus)
    8 = Landscape mode (for landscape photos with the background in focus)
    """
    exp_program = fetch("ExposureProgram", exif)
    # ISO
    iso = fetch("ISOSpeedRatings", exif)
    # GainControl ?
    """
    Indicates the degree of overall image gain adjustment.
    0 = None
    1 = Low gain up
    2 = High gain up
    3 = Low gain down
    4 = High gain down
    """
    gain = fetch("GainControl", exif)
    # summary
    info = ""
    if focal and focal_in35mm:
        info += "{}mm ({}mm)".format(focal, focal_in35mm)
    if fnumber:
        info += ", F{}".format(fnumber)
    if exp_time:
        info += ", {}s".format(exp_time)
    if iso:
        info += ", ISO {}{}".format(iso, ev)
    if model:
        info += ", {}".format(model)
    if ct:
        info += ", {}".format(ct)
    if info[0] == ",":
        info = info[1:]
    info = info.strip()
    return info, img

def transform(path, exp_ratio, verbose):
    img = PIL.Image.open(path)

    # load info
    try:
        info, img = load_info(img, path, verbose)
    except:
        print(red("Unexpected error:{}".format(sys.exc_info()[0])))
        info = os.path.basename(path)
    print(yellow("    info: {}".format(info)))

    # re-size
    w, h = img.size
    ratio = float(w) / h
    print("    size: w{}px, h{}px, ".format(w, h) + "{0:.3f}".format(ratio))
    # expectation
    if ratio < 1.0:
        exp_ratio = 1.0 / exp_ratio

    d_w = max(int(exp_ratio * h - w), 0)
    d_h = max(int(float(w) / exp_ratio - h), 0)

    attach_info = (info != None)
    info_on_bottom, info_on_right = False, False
    min_extra = 120
    if attach_info:
        info_on_right = (d_w > 0)
        info_on_bottom = (d_h > 0)
        # if perfect ratio, force to print on bottom
        if (not info_on_right) and (not info_on_bottom):
            info_on_bottom = True
        if info_on_right and d_w < min_extra:
            d_w = min_extra
            d_h = int(float(d_w + w) / exp_ratio - h)
        if info_on_bottom and d_h < min_extra:
            d_h = min_extra
            d_w = int(float(d_h + h) * exp_ratio - w)

    print("  resize: w{}px, h{}px, ".format(w + d_w, h + d_h) + "{0:.3f}".format(exp_ratio))
    # if too much margin, place image to the center
    # otherwise, place it to one side
    start_h = 0
    if d_h / 2 > min_extra:
        start_h = int(d_h / 2)
    start_w = 0
    if d_w / 2 > min_extra:
        start_w = int(d_w / 2)

    # new image
    final = Image.new(mode = 'RGB', size = (w + d_w, h + d_h), color = (0, 0, 0))
    final.paste(img, (start_w, start_h))

    # attach info
    if attach_info:
        fnt = ImageFont.truetype('./fonts/Arial.ttf', 80)
        margin = 20
        if info_on_bottom:
            txt_w, txt_h = (w, d_h - start_h)
        if info_on_right:
            txt_w, txt_h = (h, d_w - start_w)
        txt_img = Image.new('RGB', (txt_w, txt_h), (0, 0, 0))
        txt_d = ImageDraw.Draw(txt_img)
        # Note, info += " " to avoid cut text on the edge
        font_w, font_h = txt_d.textsize(info + " ", font=fnt)
        txt_d.text((txt_w - font_w - margin, margin), info, font=fnt, fill=(200, 200, 200))
        if info_on_bottom:
            final.paste(txt_img, (start_w, start_h + h))
        if info_on_right:
            txt_img = txt_img.rotate(270, expand=1)
            final.paste(txt_img, (start_w + w, start_h))

    return final, attach_info, info_on_bottom, info_on_right, d_w, d_h

def is_image(path):
    str = os.path.basename(path).lower()
    if str.startswith("."):
        return False 
    return str.endswith(".png") or str.endswith(".jpg") or str.endswith(".heic") or str.endswith(".jpeg")

def output_prefix(attach_info, info_on_bottom, info_on_right, d_w, d_h):
    p = ""
    if attach_info:
        if info_on_bottom:
            p += "b-"
        if info_on_right:
            p += "r-"
    else:
        p += "no-"
    suffix = False
    if d_w > 0:
        p += "w" + str(d_w)
        suffix = True
    if d_h > 0:
        p += "h" + str(d_h)
        suffix = True
    if suffix:
        p += "-"
    return p

def process_folder(input_dir, output_dir, size, verbose):
    count = 0
    for file in os.listdir(input_dir):
        path = os.path.join(input_dir, file)
        count += process_img(path, output_dir, size, verbose)
    print(green("Proccessed {} images".format(count)))

def process_img(path, output_dir, size, verbose):
    if not is_image(path):
        print(red("ignore {}".format(path)))
        return 0
    ratio = pick_ratio(size)
    print("processing {}".format(path))
    img, attach_info, info_on_bottom, info_on_right, d_w, d_h = transform(path, ratio, verbose)
    #img.show()
    prefix = output_prefix(attach_info, info_on_bottom, info_on_right, d_w, d_h)
    outpath = os.path.join(output_dir, prefix + os.path.basename(path))
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    img.save(outpath, "JPEG", quality=100, optimize=True, progressive=True)
    print("   saved: {}".format(outpath))
    return 1

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', dest='input',  required=True, help='input directory or file')
    parser.add_argument('-v', dest='verbose', help='verbose', required=False, action='store_true')
    parser.add_argument('-o', dest='output_dir', help='output directory')
    parser.add_argument('-s', dest='size', help='photo size: "5", "5e", "6"', default="5")
    args = parser.parse_args()

    if os.path.isdir(args.input):
        if not args.output_dir:
            args.output_dir = args.input + "-output"
        process_folder(args.input, args.output_dir, args.size, args.verbose)
    else:
        if not args.output_dir:
            args.output_dir = os.path.dirname(args.input) + "-output"
        process_img(args.input, args.output_dir, args.size, args.verbose)


