#!/usr/bin/env python3
# udt-accent: pick a Catppuccin Macchiato accent matching a wallpaper.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License v2 only.
"""Extract a wallpaper's signature color and snap it to a Macchiato accent.

Usage: udt-accent <wallpaper-path>
       udt-accent --selftest

The extraction deliberately calls pywal's colorz backend directly rather than
running `wal -i`, because `wal -i` rewrites the whole ~/.cache/wal directory
including the terminal's ANSI colors. See the design spec for why that matters.
"""

import json
import math
import os
import subprocess
import sys
import tempfile
from pathlib import Path

# The candidate accents. rosewater and flamingo are excluded as near-neutral
# tints that capture saturated inputs, and maroon and sapphire as near-duplicate
# hues of red and sky. lavender sits close to mauve but is kept: it is the
# desktop-wide fallback, used by Qt, GTK and anything else that needs one fixed
# accent, so it has to be a value snap() can also return.
ACCENTS = {
    "pink":     "#f5bde6",
    "mauve":    "#c6a0f6",
    "lavender": "#b7bdf8",
    "red":    "#ed8796",
    "peach":  "#f5a97f",
    "yellow": "#eed49f",
    "green":  "#a6da95",
    "teal":   "#8bd5ca",
    "sky":    "#91d7e3",
    "blue":   "#8aadf4",
}

FALLBACK = "lavender"
MIN_CHROMA = 10.0

OUTPUT = Path.home() / ".cache" / "wal" / "udt-accent.rasi"
BORDER_OUTPUT = Path.home() / ".cache" / "wal" / "udt-border.lua"
DUNSTRC = Path.home() / ".cache" / "wal" / "dunstrc"
COLORS_JSON = Path.home() / ".cache" / "wal" / "colors.json"

# Catppuccin Macchiato, in pywal's slot order. Firefox is themed from this
# via pywalfox, which reads colors.json and nothing else.
MACCHIATO = {
    "background": "#24273a", "foreground": "#cad3f5", "cursor": "#f4dbd6",
    "colors": [
        "#494d64", "#ed8796", "#a6da95", "#eed49f",
        "#8aadf4", "#f5bde6", "#8bd5ca", "#b8c0e0",
        "#5b6078", "#ed8796", "#a6da95", "#eed49f",
        "#8aadf4", "#f5bde6", "#8bd5ca", "#a5adcb",
    ],
}


def _to_lab(hexval):
    """Convert #rrggbb to CIELAB. sRGB D65, the standard conversion."""
    r, g, b = (int(hexval[i:i + 2], 16) / 255 for i in (1, 3, 5))

    def linear(c):
        return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4

    r, g, b = linear(r), linear(g), linear(b)

    x = (0.4124 * r + 0.3576 * g + 0.1805 * b) / 0.95047
    y = (0.2126 * r + 0.7152 * g + 0.0722 * b)
    z = (0.0193 * r + 0.1192 * g + 0.9505 * b) / 1.08883

    def f(t):
        return t ** (1 / 3) if t > 0.008856 else 7.787 * t + 16 / 116

    fx, fy, fz = f(x), f(y), f(z)
    return (116 * fy - 16, 500 * (fx - fy), 200 * (fy - fz))


def _hue(hexval):
    """Perceptual hue angle in radians."""
    _, a, b = _to_lab(hexval)
    return math.atan2(b, a)


def _chroma(hexval):
    """Distance from the neutral axis. Near-greys sit close to zero."""
    _, a, b = _to_lab(hexval)
    return math.hypot(a, b)


def snap(hexval):
    """Return the name of the nearest candidate accent by perceptual hue."""
    if _chroma(hexval) < MIN_CHROMA:
        return FALLBACK

    target = _hue(hexval)

    def distance(name):
        delta = abs(_hue(ACCENTS[name]) - target)
        return min(delta, 2 * math.pi - delta)  # hue is circular

    return min(ACCENTS, key=distance)


def signature_color(image):
    """Extract the image's most chromatic mid-tone color.

    Calls the colorz backend directly. It returns a list and writes nothing,
    which is what keeps the pywal cache (and so the terminal) untouched.
    """
    from pywal.backends import colorz

    colors = colorz.get(str(image), 16)
    # Slot 0 trends near-black and the upper slots near-white; the signature
    # color of an image lives in the middle.
    return max(colors[1:7], key=_chroma)


def write_accent(name):
    """Write the accent rasi file atomically."""
    hexval = ACCENTS[name]
    content = (
        "/* Generated by udt-accent. Do not edit. */\n"
        f"* {{ accent: {hexval}ff; }}\n"
    )

    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    # Write-then-rename: a rofi launch during a wallpaper change must never
    # read a half-written file.
    fd, tmp = tempfile.mkstemp(dir=str(OUTPUT.parent), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as handle:
            handle.write(content)
        os.replace(tmp, OUTPUT)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def hue_neighbours(name):
    """The accent either side of `name` on the perceptual hue wheel.

    The animated border rotates a gradient, so it needs more than one colour to
    show motion. Using the accent's own neighbours keeps the movement visible
    while staying inside one region of the palette.
    """
    order = sorted(ACCENTS, key=lambda n: _hue(ACCENTS[n]))
    i = order.index(name)
    return order[(i - 1) % len(order)], order[(i + 1) % len(order)]


def write_border(name):
    """Write the Hyprland border gradient, atomically.

    Emitted as Lua rather than a .conf snippet: Hyprland's Lua parser refuses
    `hyprctl keyword` ("keyword can't work with non-legacy parsers") and has no
    source directive, so the config reads this file with dofile() instead.
    """
    before, after = hue_neighbours(name)
    stops = ", ".join(f'"rgb({ACCENTS[n].lstrip("#")})"' for n in (before, name, after))

    content = (
        "-- Generated by udt-accent. Do not edit.\n"
        f"return {{ {stops} }}\n"
    )

    BORDER_OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=str(BORDER_OUTPUT.parent), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as handle:
            handle.write(content)
        os.replace(tmp, BORDER_OUTPUT)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def write_colors_json(name, image):
    """Rewrite colors.json as Macchiato with the accent in the highlight slots.

    pywalfox reads this file and nothing else, so this is how Firefox tracks the
    wallpaper. pywal wrote the file moments earlier with wallpaper-derived ANSI
    colours; this replaces them wholesale, which is the point: the terminal
    palette stays fixed Macchiato while only the accent moves.
    """
    hexval = ACCENTS[name]
    colors = list(MACCHIATO["colors"])
    # Slots 4 and 12 are pywalfox's link/highlight colour.
    colors[4] = colors[12] = hexval

    doc = {
        "wallpaper": str(image),
        "alpha": "100",
        "special": {
            "background": MACCHIATO["background"],
            "foreground": MACCHIATO["foreground"],
            "cursor": hexval,
        },
        "colors": {f"color{i}": c for i, c in enumerate(colors)},
    }

    fd, tmp = tempfile.mkstemp(dir=str(COLORS_JSON.parent), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as handle:
            json.dump(doc, handle, indent=4)
            handle.write("\n")
        os.replace(tmp, COLORS_JSON)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise

    # Firefox only picks the new colours up when pywalfox pushes them. Never
    # fatal: pywalfox may not be installed, and the desktop theme is unaffected.
    subprocess.run(["pywalfox", "update"], capture_output=True, check=False)


def write_dunst(name):
    """Substitute the accent into the dunst config pywal just rendered.

    pywal owns that file, so this runs after it and rewrites in place. A
    missing file is not an error: dunst simply may not be configured here.
    """
    try:
        text = DUNSTRC.read_text()
    except FileNotFoundError:
        return

    if "@ACCENT@" not in text:
        return

    fd, tmp = tempfile.mkstemp(dir=str(DUNSTRC.parent), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as handle:
            handle.write(text.replace("@ACCENT@", ACCENTS[name]))
        os.replace(tmp, DUNSTRC)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise

    # Restart dunst so it rereads the file. It must be started again, not just
    # killed: nothing else respawns it, and a dead dunst means no notifications
    # at all.
    if subprocess.run(["pkill", "-x", "dunst"], capture_output=True).returncode == 0:
        subprocess.Popen(["dunst"], start_new_session=True,
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def main(image):
    try:
        name = snap(signature_color(image))
    except Exception as exc:
        # A broken image must still leave a working theme.
        print(f"udt-accent: {exc}, falling back to {FALLBACK}", file=sys.stderr)
        name = FALLBACK

    write_accent(name)
    write_border(name)
    write_dunst(name)
    write_colors_json(name, image)

    # Hyprland only rereads its config on request, and may not be running.
    subprocess.run(["hyprctl", "reload"], capture_output=True, check=False)

    before, after = hue_neighbours(name)
    print(f"{name} {ACCENTS[name]} (border: {before} .. {name} .. {after})")


def selftest():
    # Every accent must snap to itself, or the metric is not self-consistent.
    for name, hexval in ACCENTS.items():
        got = snap(hexval)
        assert got == name, f"{name} ({hexval}) snapped to {got}"

    # Representative real-world inputs.
    assert snap("#ff8800") == "peach", snap("#ff8800")
    assert snap("#00cc44") == "green", snap("#00cc44")

    # A near-grey has an unstable hue angle and must take the fallback.
    assert snap("#888888") == FALLBACK, snap("#888888")
    assert FALLBACK in ACCENTS, FALLBACK

    # Border neighbours must be distinct from the accent and from each other,
    # or the gradient has nothing to animate between.
    for name in ACCENTS:
        before, after = hue_neighbours(name)
        assert len({before, name, after}) == 3, (name, before, after)
    assert hue_neighbours("peach") == ("red", "yellow"), hue_neighbours("peach")

    print("selftest OK")


if __name__ == "__main__":
    if len(sys.argv) == 2 and sys.argv[1] == "--selftest":
        selftest()
    elif len(sys.argv) == 2:
        main(sys.argv[1])
    else:
        print(__doc__, file=sys.stderr)
        sys.exit(2)