#!/usr/bin/env python3
#coding: utf-8
import cv2
import os
from os.path import join
import shutil

def pic_size_find(pic_path: str):
    img = cv2.imread(pic_path,0)
    h, w = img.shape[:2]
    return(w, h)

def pics_list(path:os.path):
    files = [f for f in os.listdir(path) if (f.endswith(".jpg") or f.endswith(".jpeg") or f.endswith(".png"))]

    return files

def file_move(file, dest):
    shutil.move(file, dest)

def delete_empty_folders(rootdir):
    folders = [f.path for f in os.scandir(rootdir) if f.is_dir()]
    for folder in folders:
        if not os.listdir(folder):
            print(f"{folder} is empty and will be deleted")
            os.rmdir(folder)


def main():
    #------------------ CONFIG START
    #Where the wallpaper folders are, must end in "/"
    rootdir = "/home/justine/NAS/Gallery/"

    #Folder names for landscape
    folder4k = "Wallpapers_4K"
    folder2k = "Wallpapers_2k"
    folderhd = "Wallpapers_hd"
    foldersd = "Wallpapers_other_def"
    folderuw = "UltrawideWallpapers"

    #Folder names for portrait
    vert4k = "Vertical_4k"
    vert2k = "Vertical_2k"
    verthd = "Vertical_hd"
    vertuw = "Vertical_ultrawide"

    #Resolutions to check against
    res4k = (4096, 2160)
    reslow4k = (3840, 2160)
    res2k = (2560, 1440)
    reshd = (1920, 1080)
    resuw = (3440, 1440)

    #Where to look for the unsorted wallpapers
    orig_folder = "/home/justine/Pictures/Unsorted"
    #------------------ CONFIG END

    print("Size begins !")

    for folder in [folder4k, folder2k, folderhd, foldersd, folderuw, vert2k, vert4k, verthd, vertuw]:
        os.makedirs(join(rootdir, folder), exist_ok=True)
        print(f"Folder {rootdir}{folder} created if needed")



    files = pics_list(orig_folder)

    print(f"About to run on {len(files)} pictures...")
    cont = input("Press enter to start, Ctrl+C to cancel")

    for f in files:
        f_path = join(orig_folder, f)
        size = pic_size_find(f_path)
        rsize = size[::-1]
        print(f"{f} => {size}")
        #Check landscape first
        if size == res4k or size == reslow4k:
            file_move(f_path, join(rootdir, folder4k, f))
            dest_folder = folder4k
        elif size == res2k:
            file_move(f_path, join(rootdir, folder2k, f))
            dest_folder = folder2k
        elif size == reshd:
            file_move(f_path, join(rootdir, folderhd, f))
            dest_folder = folderhd
        elif size == resuw:
            file_move(f_path, join(rootdir, folderuw, f))
            dest_folder = folderuw

        #Check portrait by inverting sizes
        elif rsize == res4k or rsize == reslow4k:
            file_move(f_path, join(rootdir, vert4k, f))
            dest_folder = vert4k
        elif rsize == res2k[::-1]:
            file_move(f_path, join(rootdir, vert2k, f))
            dest_folder = vert2k
        elif rsize == reshd[::-1]:
            file_move(f_path, join(rootdir, verthd, f))
            dest_folder = verthd
        elif rsize == resuw[::-1]:
            file_move(f_path, join(rootdir, vertuw, f))
            dest_folder = vertuw

        #...Or put in "SD" folder
        else:
            file_move(f_path, join(rootdir, foldersd, f))
            dest_folder = foldersd
        print(f"{f} moved to {rootdir}{dest_folder}")

    delete_empty_folders(rootdir)

    print("Done, Sayonara")

if __name__ == "__main__":
    main()