#!/usr/bin/python3

import argparse
from typing import List
from pathlib import Path

from PIL import Image


SOURCES = ["konachan", "wallhaven", "yandere", "moe_imouto"]
RATIO_PATHS = ("4_3", "16_9", "16_10", "vertical")


def get_wallpaper_source(filename: Path) -> str:

    if filename.name.startswith("Konachan.com"):
        return "konachan"
    elif filename.name.startswith("wallhaven"):
        return "wallhaven"
    elif filename.name.startswith("yande.re"):
        return "yandere"
    elif filename.name.startswith("moe"):
        return "moe_imouto"
    else:
        return ""


def ratio_and_resolution(filename: Path, guess_source=True) -> Path:

    img = Image.open(filename)
    resolution = f"{img.width}x{img.height}"
    ratio = img.width / img.height

    check_ratios = [("4:3", 4/3), ("16:9", 16/9), ("16:10", 16/10)]
    # Calculate the distance from each ratio of the image in question
    distances = {label: (ratio-ref)**2 for label, ref in check_ratios}
    # Calculate the ratio with the minimum distance
    min_dist = min(distances, key=distances.get).replace(":", "_")

    if guess_source:
        source = get_wallpaper_source(filename)
    else:
        source = ""

    new_filename = "_".join((filename.stem, resolution)) + filename.suffix

    # Special case phone wallpapers for now
    if img.height > img.width:
        min_dist = "vertical"

    if source:
        final_name = Path(min_dist) / source / new_filename
    else:
        final_name = Path(min_dist) / new_filename

    return final_name


def organize_images(files: List[Path],
                    destination_path: Path, dry_run=True) -> None:

    for ratio_set in RATIO_PATHS:
        if not dry_run:
            (destination_path / ratio_set).mkdir(exist_ok=True)
            for source in SOURCES:
                (destination_path / ratio_set / source).mkdir(exist_ok=True)

    for element in files:
        new_name = destination_path / ratio_and_resolution(element)

        if dry_run:
            print(f"{element} -> {new_name}")
            continue

        if not new_name.exists():
            element.rename(new_name)
            # new_name.symlink_to(element)


def get_image_list(source: Path) -> List[Path]:

    candidates = list()
    for element in source.glob("**/*"):

        if element.is_dir():
            continue

        # Exclude already processed directories
        if any(part in RATIO_PATHS for part in element.parts):
            continue

        if element.suffix in (".png", ".jpg", ".webp", ".jpeg", ".bmp"):
            candidates.append(element)

    return candidates


def cleanup(destination: Path) -> None:

    for element in RATIO_PATHS:
        full_path = destination / element

        for source in SOURCES:
            source_path = full_path / source

            if not any(source_path.iterdir()):
                source_path.rmdir()

        if not any(full_path.iterdir()):
            full_path.rmdir()


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true",
                        help="Only simulate")
    parser.add_argument("source", help="Source wallpaper directory")
    parser.add_argument("destination", help="Destination path to move data to")

    options = parser.parse_args()
    source = Path(options.source).absolute()

    destination = Path(options.destination).absolute()

    images = get_image_list(source)
    organize_images(images, destination, options.dry_run)

    cleanup(destination)


if __name__ == "__main__":
    main()