dotfiles/scripts/lib/wallpaper.py

"""Script to update the wallpaper"""
from abc import ABC, abstractmethod
from json import dumps, loads
from math import ceil, floor, isnan, nan
from os import environ
from pathlib import Path
from random import sample
from re import Pattern, compile, match
from subprocess import run
from time import sleep
from typing import cast, override

from .evening import on_steam, sunup_amount


# Python sets aren't insertion-ordered, but dict keys are
type OrderedSet[T] = dict[T, None]


class Background(ABC):
	_colour: Pattern[str] = compile(r"^[a-fA-F0-9]{6}$")

	@abstractmethod
	def blue_light(self) -> float: ...

	@abstractmethod
	@override
	def __eq__(self, value: object, /) -> bool: ...

	@abstractmethod
	@override
	def __hash__(self) -> int: ...

	@abstractmethod
	@override
	def __str__(self) -> str: ...

	@classmethod
	def parse(cls, arg: str) -> "Background":
		if match(cls._colour, arg):
			return Colour(int(arg, 16))
		else:
			return Wallpaper(Path(arg))


class Wallpaper(Background):
	_cache: Path = Path(environ["XDG_CACHE_HOME"]) / "wallpaper_blue_light.json"
	_wallpapers: Path = Path("~/Pictures/wallpapers/").expanduser()

	def __init__(self, image: Path) -> None:
		self._path: Path = image
		self._blue_light: float = nan

	def _calculate_blue_light(self) -> float:
		from PIL import Image
		data = cast(list[int], Image.open(self._path).getdata(2))
		return sum(data) / len(data)

	@classmethod
	def _read_cache(cls) -> dict[str, float]:
		if cls._cache.exists():
			if (text := cls._cache.read_text().strip()):
				if isinstance((obj := loads(text)), dict):
					return cast(dict[str, float], obj)
		return {}

	@classmethod
	def _write_cache(cls, cache: dict[str, float]) -> None:
		stringified = dumps(cache)
		if len(stringified) != cls._cache.write_text(stringified):
			raise Exception("Did not fully write blue light cache to file")

	@classmethod
	def available(cls) -> "OrderedSet[Background]":
		cache = cls._read_cache()
		modified = False
		wallpapers = {str(path): Wallpaper(path) for path in cls._wallpapers.iterdir()}
		for wallpaper in wallpapers.values():
			cached = cache.get(str(wallpaper._path))
			if cached is None or isnan(cached):
				print("New file:", wallpaper._path)
				cached = wallpaper._calculate_blue_light()
				modified = True
				cache[str(wallpaper._path)] = cached
			wallpaper._blue_light = cached
		if modified:
			cache = {path: score for (path, score) in sorted(cache.items(), key=lambda pair: pair[1])}
			cls._write_cache(cache)
		return {wallpapers[path]: None for path in cache.keys()}

	@override
	def blue_light(self) -> float:
		return self._blue_light

	@override
	def __eq__(self, value: object, /) -> bool:
		if isinstance(value, Wallpaper):
			return self._path == value._path
		return False

	@override
	def __hash__(self) -> int:
		return hash(self._path)

	@override
	def __str__(self) -> str:
		return str(self._path)

	@override
	def __repr__(self) -> str:
		return f"Wallpaper({repr(self._path)})"


class Colour(Background):
	def __init__(self, value: int) -> None:
		self._value: int = value

	@override
	def blue_light(self) -> float:
		return self._value & 0xFF

	@override
	def __eq__(self, value: object, /) -> bool:
		if isinstance(value, Colour):
			return self._value == value._value
		return False

	@override
	def __hash__(self) -> int:
		return hash(self._value)

	@override
	def __str__(self) -> str:
		return hex(self._value)

	@override
	def __repr__(self) -> str:
		return f"Colour({hex(self._value)})"


def _convert_bytes(output: bytes) -> str:
	return output.decode("utf-8").strip()


def _swww_cmd(*args: str) -> str:
	output = run(["swww", *args], capture_output=True)
	if output.stderr:
		raise Exception(f"swww returned: '{_convert_bytes(output.stderr)}'")
	return _convert_bytes(output.stdout)


class Monitor:
	_query: Pattern[str] = compile(r"^([^:]+): (\d+)x(\d+), scale: ([\d.]+), currently displaying: (?:image|color): (.+)$")
	_transition_args: list[str] = [
		"--transition-type=wipe",
		"--transition-angle=30",
		"--transition-step=45",
		"--transition-fps=120",
	]

	def __init__(self, line: str) -> None:
		if not (values := match(self._query, line)):
			raise Exception(f"Failed to extract information from: '{line}'")
		self._name: str = values[1]
		self._resolution: tuple[int, int] = (int(values[2]), int(values[3]))
		self._scale: float = float(values[4])
		self._background: Background = Background.parse(values[5])

	def is_colour(self) -> bool:
		return isinstance(self._background, Colour)

	def set_wallpaper(self, background: Background) -> None:
		self._background = background
		_ = _swww_cmd("img", *self._transition_args, f"--outputs={self._name}", str(self._background))

	def get_background(self) -> Background:
		return self._background

	@classmethod
	def current(cls) -> "list[Monitor]":
		return [Monitor(line) for line in _swww_cmd("query").splitlines()]

	@override
	def __eq__(self, value: object, /) -> bool:
		if isinstance(value, Monitor):
			return self._background == value._background
		return False

	@override
	def __repr__(self) -> str:
		return f"Monitor('{self._name}: {self._resolution[0]}x{self._resolution[1]}, scale: {self._scale}, currently displaying: {"color" if self.is_colour() else "image"}: {self._background}')"

	@override
	def __str__(self) -> str:
		return f"{self._name}: {self._background}"


def select_wallpapers(available: list[Background]) -> list[Background]:
	radius = len(available) / 3
	centre = sunup_amount() * len(available)
	return available[max(0, floor(centre - radius)): min(len(available) + 1, ceil(centre + radius))]


def main(check_steam: bool = False) -> None:
	if check_steam and on_steam():
		return

	monitors = Monitor.current()
	wallpapers = Wallpaper.available()

	for monitor in monitors:
		wallpapers.pop(monitor.get_background(), None)

	for monitor, wallpaper in zip(monitors, sample(select_wallpapers(list(wallpapers.keys())), len(monitors))):
		monitor.set_wallpaper(wallpaper)
		sleep(1.25)