#!/usr/bin/env python3

# Description:
#   Change color theme for various applications based on an image, color, or user selection.
#   This is done by running application-specific scripts located in the config directory.
#
# Requirements:
#   - colorthief (python3 package) # too lazy to implement color extraction myself :D

import os
import sys
from pathlib import Path
import argparse

PALETTES = {
    "catppuccin-mocha": {
        "rosewater": "f5e0dc",
        "flamingo": "f2cdcd",
        "pink": "f5c2e7",
        "mauve": "cba6f7",
        "red": "f38ba8",
        "maroon": "eba0ac",
        "peach": "fab387",
        "yellow": "f9e2af",
        "green": "a6e3a1",
        "teal": "94e2d5",
        "sky": "89dceb",
        "sapphire": "74c7ec",
        "blue": "89b4fa",
        "lavender": "b4befe",
    },
}

CONFIG_DIR = Path("~/.config").expanduser()

# An application may have multiple scripts (e.g. due to config-switch)
SCRIPTS = {
    "eww": [CONFIG_DIR / "eww" / "apply-color"],
    "fastfetch": [CONFIG_DIR / "fastfetch" / "apply-color"],
    "fuzzel": [CONFIG_DIR / "fuzzel" / "apply-color"],
    "hypr": [CONFIG_DIR / "hypr" / "apply-color"],
    "kvantum": [CONFIG_DIR / "fish" / "apply-color-kvantum"],  # borrowing fish's directory
    "nwg-look": [CONFIG_DIR / "fish" / "apply-color-nwg-look"],  # borrowing fish's directory
    "mako": [CONFIG_DIR / "mako" / "apply-color"],
    "niri": [CONFIG_DIR / "niri" / "apply-color"],
    "oh-my-posh": [CONFIG_DIR / "fish" / "apply-color-omp"],  # borrowing fish's directory
    "quickshell": [CONFIG_DIR / "quickshell" / "apply-color"],
    "rofi": [CONFIG_DIR / "rofi" / "apply-color"],
    "waybar": [CONFIG_DIR / "waybar" / "apply-color"],
    "wlogout": [CONFIG_DIR / "wlogout" / "apply-color", CONFIG_DIR / "wlogout-niri" / "apply-color"],
    "yazi": [CONFIG_DIR / "yazi" / "apply-color"],
}
# or simply `find ${CONFIG_DIR} -type f -iname 'apply-color*'` to get all available scripts,
# but I need the exact application names anyway, so hardcoding does make some sense


def hex2rgb(hex_color: str) -> tuple[int, int, int]:
    return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))  # type: ignore


def extract_color(image_path: str) -> str:
    from colorthief import ColorThief
    return "#{:02x}{:02x}{:02x}".format(*ColorThief(image_path).get_color(quality=10))


def match_color(color: str, palette: dict[str, str]) -> str:
    """ Matches a given color (rrggbb hex) to the closest color in the palette."""
    # HUE distance of the given and returned color must no<t exceed this value
    HUE_THRESHOLD = 60.0  # degrees

    color = color.lower().strip().removeprefix('#')

    # weigh by CCIR 601 luminosity
    fr, fg, fb = 0.299 / 255 / 255, 0.587 / 255 / 255, 0.114 / 255 / 255
    lfr, lfg, lfb = 0.299 / 255, 0.587 / 255, 0.114 / 255

    def color_distance(c1: str, c2: str) -> float:
        r1, g1, b1 = hex2rgb(c1)
        r2, g2, b2 = hex2rgb(c2)
        diff_l = (lfr * (r1 - r2) + lfg * (g1 - g2) + lfb * (b1 - b2))
        diff_r = fr * (r1 - r2) ** 2
        diff_g = fg * (g1 - g2) ** 2
        diff_b = fb * (b1 - b2) ** 2
        return (diff_r + diff_g + diff_b) * 0.75 + diff_l ** 2

    def color_distance_hue(c1: str, c2: str) -> float:
        def rgb2hue(r, g, b) -> float:
            r, g, b = r / 255.0, g / 255.0, b / 255.0
            mx = max(r, g, b)
            mn = min(r, g, b)
            diff = mx - mn

            if diff == 0:
                return 0.0

            if mx == r:
                hue = (g - b) / diff + (6 if g < b else 0)
            elif mx == g:
                hue = (b - r) / diff + 2
            else:
                hue = (r - g) / diff + 4

            return hue * 60
        r1, g1, b1 = hex2rgb(c1)
        r2, g2, b2 = hex2rgb(c2)
        return abs(rgb2hue(r1, g1, b1) - rgb2hue(r2, g2, b2))

    closest_color = min(palette.keys(), key=lambda k: color_distance(color, palette[k]))
    print(f"Matched color {color} to {closest_color}")

    # if the hue distance is too large, rematch
    if color_distance_hue(color, palette[closest_color]) > HUE_THRESHOLD:
        print(f"Color {color} is too far from {closest_color}, rematching'")
    else:
        return closest_color

    closest_color = min(palette.keys(), key=lambda k: color_distance_hue(color, palette[k]))
    print(f"Rematched color {color} to {closest_color}")

    return closest_color


def pick_flavor(palette: dict[str, str]) -> str:
    def is_interactive() -> bool:
        return sys.stdin.isatty() and sys.stdout.isatty()

    def is_truecolor() -> bool:
        colorterm = os.environ.get('COLORTERM', '')
        term = os.environ.get('TERM', '')

        return (
            'truecolor' in colorterm or
            '24bit' in colorterm or
            term.endswith('-256color')
        )

    if is_interactive():
        isTruecolor = is_truecolor()
        print("Available flavors:")
        for i, flavor in enumerate(palette.keys(), 1):
            r, g, b = hex2rgb(palette[flavor])
            if isTruecolor:
                print(f"\033[38;2;{r};{g};{b}m█ {i}. {flavor}: #{palette[flavor]}\033[0m")
            else:
                print(f"{i}. {flavor}")
        while True:
            choice = input("Pick a flavor by number: ")
            if choice.isdigit() and 1 <= int(choice) <= len(palette):
                return list(palette.keys())[int(choice) - 1]
            print("Invalid choice. Try again.")
    else:
        print("No flavor specified.")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Change color theme for various applications.")
    parser.add_argument('-i', '--image', type=str, help="Path to the image")
    parser.add_argument('-f', '--flavor', type=str, help="Flavor to apply")
    parser.add_argument('-c', '--color', type=str, help="Color to match from the palette")
    parser.add_argument('arguments', nargs='*',
                        help="'app1 !app2' to include(only) / exclude(all but) specific applications. "
                        "Available apps: " + ', '.join(SCRIPTS.keys()))

    arguments = parser.parse_args()

    # for future use, probably
    def parse_palette_name() -> str:
        return "catppuccin-mocha"

    def parse_flavor(palette: dict[str, str]) -> str:
        if arguments.flavor:
            if arguments.flavor not in palette:
                print(f"Unknown flavor: {arguments.flavor}. Available flavors: {', '.join(palette.keys())}")
                sys.exit(1)
            flavor = arguments.flavor
        elif arguments.color:
            flavor = match_color(arguments.color, palette)
            print(f"Matched color: {flavor}")
        elif arguments.image:
            if not Path(arguments.image).exists():
                print(f"Image file {arguments.image} does not exist.")
                sys.exit(1)
            color = extract_color(arguments.image)
            print(f"Extracted color {color} from image {arguments.image}")
            flavor = match_color(color, palette)
            print(f"Matched color: {flavor}")
        else:
            flavor = pick_flavor(palette)
        return flavor

    def parse_apps() -> tuple[set[str], set[str]]:
        includes = set()
        excludes = set()
        for arg in arguments.arguments:
            if arg.startswith('!'):
                excludes.add(arg[1:])
            else:
                includes.add(arg)
        return includes, excludes

    palette_name = parse_palette_name()
    palette = PALETTES[palette_name]
    flavor = parse_flavor(palette)
    includes, excludes = parse_apps()

    apps = set()
    if includes:
        print(f"Including only: {', '.join(includes)}")
        for app in includes:
            if app in SCRIPTS:
                apps.add(app)
            else:
                print(f"Unknown application: {app}. Available applications: {', '.join(SCRIPTS.keys())}")
                sys.exit(1)
    else:
        apps = set(SCRIPTS.keys())

    if excludes:
        print(f"Excluding: {', '.join(excludes)}")
        apps -= excludes

    print(f"Applying flavor '{flavor}' for {len(apps)} applications.")

    for app in apps:
        for script in SCRIPTS[app]:
            print(f"Running {script}:")
            ret = os.system(f'"{script}" {palette_name} {flavor} {palette[flavor]}')
            print(f"{script} exited with code {ret}")
            print("")

    os.system(
        f'notify-send -a "change-colortheme" "Colortheme Changed" "Palette: {palette_name};\nFlavor: {flavor};\nApplied to {len(apps)} applications."')


if __name__ == "__main__":
    main()
