#!/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 re
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"
QML_OUTPUT = Path.home() / ".cache" / "wal" / "udt-palette.qml"

# The installed palette, which is a symlink back into this repo. Read rather
# than duplicated: palette.rasi is the one place the structural colours are
# written down, and a second copy here would drift the first time one changed.
PALETTE_RASI = Path.home() / ".config" / "rofi" / "udt" / "palette.rasi"
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 write_atomic(path, content):
    """Write a file atomically, so no reader ever sees it half-written.

    Every generated file is watched by something: rofi rereads on launch,
    quickshell watches with a FileView. A torn read shows up as a theme that
    briefly loses its colours.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as handle:
            handle.write(content)
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


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"
    )

    # Write-then-rename: a rofi launch during a wallpaper change must never
    # read a half-written file.
    write_atomic(OUTPUT, content)


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"
    )

    write_atomic(BORDER_OUTPUT, content)


def read_palette():
    """Parse palette.rasi into {name: "#rrggbb"}.

    rofi writes colours as #rrggbbaa; QML wants #rrggbb, and every entry in
    that file is fully opaque, so the alpha pair is dropped rather than
    converted. Returns an empty dict if the file is missing, which leaves the
    QML palette to fall back to its own defaults.
    """
    try:
        text = PALETTE_RASI.read_text()
    except OSError:
        return {}

    found = {}
    for key, value in re.findall(r"(\w+):\s*#([0-9a-fA-F]{6,8})\s*;", text):
        found[key] = "#" + value[:6]
    return found


def write_qml(name):
    """Write the palette and current accent as a QML singleton.

    Emitted as QML rather than parsed from palette.rasi by quickshell itself:
    the accent has to reach it anyway, so one generated file carrying both
    means a component needs a single FileView and no rasi parser. It is
    regenerated on every wallpaper change along with the other outputs.
    """
    palette = read_palette()
    if not palette:
        print("udt-accent: palette.rasi unreadable, skipping QML palette",
              file=sys.stderr)
        return

    rows = "\n".join(
        f'    readonly property color {key}: "{value}"'
        for key, value in sorted(palette.items())
    )

    content = (
        "// Generated by udt-accent. Do not edit.\n"
        "//\n"
        "// The Catppuccin Macchiato palette from palette.rasi, plus the accent\n"
        "// currently snapped from the wallpaper. Import it from a quickshell\n"
        "// component and watch this file to follow theme changes.\n"
        "pragma Singleton\n"
        "\n"
        "import QtQuick\n"
        "\n"
        "QtObject {\n"
        f'    readonly property color accent: "{ACCENTS[name]}"\n'
        f'    readonly property string accentName: "{name}"\n'
        "\n"
        f"{rows}\n"
        "}\n"
    )

    write_atomic(QML_OUTPUT, content)


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 = {
        # Resolved, not as passed: wallp gives the real file but a caller may
        # give ~/.cache/wal/wpaper, the symlink to it. Same wallpaper either
        # way, so record one spelling.
        "wallpaper": os.path.realpath(image),
        "alpha": "100",
        "special": {
            "background": MACCHIATO["background"],
            "foreground": MACCHIATO["foreground"],
            "cursor": hexval,
        },
        "colors": {f"color{i}": c for i, c in enumerate(colors)},
    }

    write_atomic(COLORS_JSON, json.dumps(doc, indent=4) + "\n")

    # 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

    write_atomic(DUNSTRC, text.replace("@ACCENT@", ACCENTS[name]))

    # 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_qml(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)
