#!/usr/bin/env python

# This simple Python 3 script simply generates a color scheme in wpgtk style with the auto-adjusted colors.
# I like how wpgtk generates the color palette but I don't like the workflow of it.
# I like the simple workflow of pywal but I don't like the colors generated by it and I think it has built-in limitations for generating colors.
# So I combined both of the best in this one glued script.

# The command line program simply needs an image path as the argument — `export-theme ~/Pictures/mountain.jpg`.
# It will also export the theme as a JSON file in the current directory — `mountain.json`.

# Most of the code are "borrowed" from the wpgtk codebase and I've simply studied them and added some documentation.
# It simply needs pywal as a dependency.

import argparse
from colorsys import rgb_to_hls, hls_to_rgb
import json
import math
from operator import itemgetter
import os.path
from pathlib import Path
import sys

import pywal


###################
# COLOR FUNCTIONS #
###################

def hls_to_hex(hls):
    """
    Returns a HLS coordinate into its hex color code equivalent.

    :param: hls - An HLS tuple according to the colorsys library (https://docs.python.org/3.8/library/colorsys.html).

    :returns: A hex color string equivalent.
    """
    h, l, s = hls
    r, g, b = hls_to_rgb(h, l, s)
    rgb_int = [max(min(int(elem), 255), 0) for elem in [r, g, b]]

    return pywal.util.rgb_to_hex(rgb_int)


def hex_to_hls(hex_string):
    """
    Returns an HLS coordinate equivalent of the given hex color code.
    It uses RGB as the intermediate for converting the hex string.

    :param: hex_string - A (hopefully) valid hex string.

    :returns: An HLS tuple compatible with the colorsys library.
    """
    r, g, b = pywal.util.hex_to_rgb(hex_string)
    return rgb_to_hls(r, g, b)


def get_distance(hex_src, hex_tgt):
    """
    Returns the distance between two hex values.
    The formula used in this function is based from the Wikipedia article (https://en.wikipedia.org/wiki/Color_difference).

    :param: hex_src - A hex color string.
    :param: hex_tgt - The target hex color code.

    :returns: A float that describes the distance.
    """
    r1, g1, b1 = pywal.util.hex_to_rgb(hex_src)
    r2, g2, b2 = pywal.util.hex_to_rgb(hex_tgt)

    return math.sqrt((r2 - r1)**2 + (g2 - g1)**2 + (b2 - b1)**2)


def alter_brightness(hex_string, light, sat=0):
    """
    Alters amount of light and saturation in a color.

    :param: hex_string - A hex color string (top-notch documentation right here, folks).
    :param: light - The amount of light to apply; generally, the acceptable range is -255 to 255 and beyond that is dangerous territory.
    :param: sat - The amount of saturation to apply; the acceptable range is -1 to 1.

    :returns: A hex color code of the adjusted color.
    """
    h, l, s = hex_to_hls(hex_string)
    l = max(min(l + light, 255), 1)
    s = min(max(s - sat, -1), 0)

    return hls_to_hex((h, l, s))


def is_dark_theme(color_list):
    """
    Checks by the color list if it's a dark theme.

    :param: color_list - A list of hex color codes usually 16 of them.

    :returns: A boolean indicating if it's a dark theme.
    """
    *_, bg_brightness = hex_to_hls(color_list[0])
    *_, fg_brightness = hex_to_hls(color_list[7])

    return fg_brightness < bg_brightness


def adjust_colors(colors, light_mode = False):
    """
    Create a clear foreground and background set of colors.

    :param: colors - A Pywal color object. See https://github.com/dylanaraps/pywal/wiki/Using-%60pywal%60-as-a-module#color-dict-format for more info.
    :param: light_mode - Toggle to create a light mode version of the adjustment.

    :returns: A Pywal color object with the adjusted values.
    """
    colors = smart_sort(colors)

    added_sat = 0.25 if light_mode else 0.1
    sign = -1 if light_mode else 1

    if light_mode == is_dark_theme(colors):
        colors[7], colors[0] = colors[0], colors[7]

    comment = [alter_brightness(colors[0], sign * 25)]
    fg = [alter_brightness(colors[7], sign * 60)]

    colors = colors[:8] + comment \
        + [alter_brightness(color, sign * hex_to_hls(color)[1] * 0.3, added_sat) for color in colors[1:7]] + fg

    return colors


def smart_sort(colors):
    """
    Automatically set the most look-alike colors to their
    corresponding place in the standard xterm colors.

    :param: colors - A list of hex color strings.

    :returns: The color list sorted out.
    """
    colors = colors[:8]
    sorted_by_color = list()
    base_colors = ["#000000", "#ff0000", "#00ff00", "#ffff00",
                   "#0000ff", "#ff00ff", "#00ffff", "#ffffff"]

    for y in base_colors:
        cd_tuple = [(x, get_distance(x, y)) for i, x in enumerate(colors)]
        cd_tuple.sort(key=itemgetter(1))
        sorted_by_color.append(cd_tuple)

    i = 0
    while i < 8:
        current_cd = sorted_by_color[i][0]
        closest_cds = [sorted_by_color[x][0] for x in range(8)]
        reps = [x for x in range(8) if closest_cds[x][0] == current_cd[0]]

        if len(reps) > 1:
            closest = min([closest_cds[x] for x in reps], key=itemgetter(1))
            reps = [x for x in reps if x != closest_cds.index(closest)]
            any(sorted_by_color[x].pop(0) for x in reps)
            i = 0
        else:
            i += 1

    sorted_colors = [sorted_by_color[x][0][0] for x in range(8)]
    return [*sorted_colors, *sorted_colors]


##################
# MAIN FUNCTIONS #
##################

def setup_args():
    """
    Setup the argument parser.
    """
    description = "A simple Pywal theme export script with auto-adjusted colors from wpgtk"
    argparser = argparse.ArgumentParser(description=description)

    argparser.add_argument("input", help="The input (image) of the color scheme to be generated.", metavar="IMAGE")
    argparser.add_argument("-o", "--output", help="The location of the colorscheme JSON to be exported.", metavar="FILE")
    argparser.add_argument("--no-output", help="Specifies no JSON output to be created; also overrides any output-related options.", action='store_true')

    return argparser


def export_wpgtk_colors(image_path):
    """
    Export Pywal templates with the given image and the color scheme data.
    Take note the exported Pywal object has no 'wallpaper' key for portability.

    :param: image_path - The path of the image. ;)

    :returns: The Pywal dictionary.
    """
    pywal_dict = pywal.colors.get(image_path)
    colors = []
    for color_name, color in pywal_dict["colors"].items():
        colors.append(color)

    colors = adjust_colors(colors)

    for index, color in enumerate(colors):
        pywal_dict["colors"][f"color{index}"] = color

    # Export every templates in Pywal including the user templates and reload the newly applied theme.
    pywal.export.every(pywal_dict)

    del pywal_dict["wallpaper"]

    # Feel free to add some reloading code that Pywal provides or something.
    # I'm not putting mine since it's intended to only export theme.
    # You can't tell me what to do here. >:-)

    return pywal_dict


def parse_args(parser, argv):
    """
    Parse the args and do the thing.

    :param: parser - An `argparse.ArgumentParser` instance.
    :param: argv - A list of arguments to be parsed.
    """
    args = parser.parse_args(argv)

    pywal_dict = export_wpgtk_colors(args.input)

    if args.output is not None:
        output_file = args.output
    else:
        output_file = f"{Path(args.input).stem}.json"

    if not args.no_output:
        # Save the adjusted theme as a JSON file.
        with open(output_file, "w") as json_file:
            json.dump(pywal_dict, json_file, indent = 4)


if __name__ == "__main__":
    argparser = setup_args()

    parse_args(argparser, sys.argv[1:])