#!/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
(which this no longer writes to, but pywal still would)
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 colour tables are generated: udt-palette renders them from
# palette/<scheme>.conf and palette/roles-<scheme>.conf, so the accents this
# snaps to follow whatever scheme is selected. Regenerate with ./install.sh.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from udt_colors import ACCENTS, FALLBACK, PALETTE, SCHEME  # noqa: E402

MIN_CHROMA = 10.0

CACHE = Path.home() / ".cache" / "udt"
OUTPUT = CACHE / "udt-accent.rasi"
BORDER_OUTPUT = CACHE / "udt-border.lua"
SUBLIME = (Path.home() / ".config" / "sublime-text" / "Packages" / "User"
           / "udt.sublime-color-scheme")
QML_OUTPUT = CACHE / "udt-palette.qml"

# The installed rofi palette, itself generated by udt-palette. Read rather than
# duplicated so the QML singleton carries exactly the colours rofi uses, under
# the same names.
PALETTE_RASI = Path.home() / ".config" / "rofi" / "udt" / "palette.rasi"
# pywalfox hardcodes <cache>/wal/colors.json and honours only XDG_CACHE_HOME,
# which would move every cache, so this one file is written where it looks.
# install.sh symlinks ~/.cache/wal/colors.json here, so the real file lives
# with the rest of the palette and ~/.cache/wal holds nothing but that link.
COLORS_JSON = CACHE / "colors.json"

# The SDDM theme's live palette. ~/.cache is 0700 and unreachable to the sddm
# user, ~/.local/share is 0755, so the file has to live here. Root symlinks
# /usr/share/sddm/themes/udt/theme.conf.user to it. See the sddm-theme-udt spec.
LIBEXEC = Path.home() / ".local" / "share" / "udt"
SDDM_OUTPUT = LIBEXEC / "sddm-theme.conf"

# The keys the theme asks for. The names are the same Catppuccin-compat names
# palette.rasi and udt-palette.qml use, so one palette has one vocabulary.
SDDM_KEYS = ["base", "mantle", "crust", "surface0", "surface1", "surface2",
             "text", "subtext0", "subtext1", "overlay0", "overlay1", "overlay2",
             "red", "yellow", "green", "teal", "blue", "lavender"]


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"
        f"// The {SCHEME} 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(PALETTE["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/udt/wpaper, the symlink to it. Same wallpaper either
        # way, so record one spelling.
        "wallpaper": os.path.realpath(image),
        "alpha": "100",
        "special": {
            "background": PALETTE["background"],
            "foreground": PALETTE["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_sublime(name):
    """Substitute the accent into the Sublime colour scheme install.sh copied.

    In place: Sublime colour schemes do not cascade, so the accent cannot
    arrive as a second file. Sublime watches Packages/User and reloads a
    scheme when its file changes, so no signal is needed. A missing file is not
    an error, Sublime simply may not be installed here.

    This reads the installed copy and has nowhere to recover @ACCENT@ from
    once substituted, so the previous accent is replaced by matching the
    variable line rather than a placeholder.
    """
    try:
        text = SUBLIME.read_text()
    except FileNotFoundError:
        return

    new, count = re.subn(r'("accent":\s*)"[^"]*"',
                         lambda m: f'{m.group(1)}"{ACCENTS[name]}"', text, count=1)
    if count and new != text:
        write_atomic(SUBLIME, new)


def sddm_conf(palette, name, image):
    """The SDDM theme config, as an INI fragment.

    theme.conf.user overrides theme.conf and is exposed to the theme as the
    `config` object. It is data parsed by QSettings, never executed, so nothing
    the login user can write reaches the greeter as code.
    """
    lines = [
        "# Generated by udt-accent. Do not edit.",
        "[General]",
        "type=image",
        f"background={os.path.realpath(image)}",
        f"accent={ACCENTS[name]}",
    ]
    lines += [f"{key}={palette.get(key, '#000000')}" for key in SDDM_KEYS]
    return "\n".join(lines) + "\n"


def write_sddm(name, image):
    """Write the live SDDM palette, atomically, world-readable.

    World-readable because the greeter runs as the sddm user. The atomic writer
    creates the file 0600, so the mode is fixed up here.
    """
    palette = read_palette()
    if not palette:
        print("udt-accent: palette.rasi unreadable, skipping SDDM theme",
              file=sys.stderr)
        return
    write_atomic(SDDM_OUTPUT, sddm_conf(palette, name, image))
    os.chmod(SDDM_OUTPUT, 0o644)


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_sublime(name)
    write_sddm(name, image)
    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. Asserted by hue rather than by name:
    # the accent table is generated, so the colour an orange input snaps to is
    # called "peach" under Catppuccin and "orange" under Tokyo Night.
    for probe in ("#ff8800", "#00cc44", "#3366ff"):
        got = snap(probe)
        delta = abs(_hue(ACCENTS[got]) - _hue(probe))
        delta = min(delta, 2 * math.pi - delta)
        assert delta < 0.6, f"{probe} snapped to {got}, {delta:.2f} rad away"

    # The Sublime accent is replaced by matching its current value, not a
    # placeholder, so it has to survive being run twice: once over the freshly
    # installed @ACCENT@, and again over the accent it wrote last time.
    probe = '{\n    "variables":\n    {\n        "accent":   "@ACCENT@",\n'
    first, n1 = re.subn(r'("accent":\s*)"[^"]*"', r'\1"#f5a97f"', probe, count=1)
    second, n2 = re.subn(r'("accent":\s*)"[^"]*"', r'\1"#8bd5ca"', first, count=1)
    third, n3 = re.subn(r'("accent":\s*)"[^"]*"', r'\1"#8bd5ca"', second, count=1)
    assert (n1, n2, n3) == (1, 1, 1), (n1, n2, n3)
    assert "@ACCENT@" not in first and '"#f5a97f"' in first
    assert '"#f5a97f"' not in second and '"#8bd5ca"' in second
    assert third == second, "accent substitution is not idempotent"

    # 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)
    # Neighbours are the adjacent hues on the wheel, so each sits nearer to
    # the accent than the accent's opposite does.
    for name in ACCENTS:
        before, after = hue_neighbours(name)
        for neighbour in (before, after):
            delta = abs(_hue(ACCENTS[neighbour]) - _hue(ACCENTS[name]))
            assert min(delta, 2 * math.pi - delta) < math.pi, (name, neighbour)

    # The SDDM theme reads this file through a symlink, so it has to carry every
    # key the theme asks for, and a second run has to overwrite cleanly rather
    # than append.
    palette = {k: "#1e2030" for k in SDDM_KEYS}
    probe_image = Path(tempfile.mkdtemp()) / "wall.png"
    probe_image.write_bytes(b"")
    text = sddm_conf(palette, FALLBACK, str(probe_image))
    for key in ["type", "background", "accent", *SDDM_KEYS]:
        assert f"\n{key}=" in "\n" + text, f"sddm conf missing {key}"
    assert text.count("accent=") == 1
    assert "background=" + str(probe_image) in text

    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)