z-bar-qt/cli/src/zshell/subcommands/battery.py

from __future__ import annotations

import contextlib
import fcntl
import json
import logging
import os
import re
import signal
import struct
import tempfile
import threading
import time
from collections.abc import Callable
from dataclasses import asdict, dataclass
from pathlib import Path

import typer
from logitech_receiver import hidpp10, hidpp20
from logitech_receiver.common import Battery, BatteryStatus, Notification
from logitech_receiver.hidpp10_constants import Registers
from logitech_receiver.hidpp20_constants import SupportedFeature

app = typer.Typer(
	help="Read live battery status from Logitech HID++ peripherals."
)

logger = logging.getLogger("zshell.battery")

DEFAULT_OUTPUT = Path.home() / ".cache" / "zshell" / "battery.json"

_RUNTIME_DIR = Path(
	os.environ.get("XDG_RUNTIME_DIR") or (Path.home() / ".cache" / "zshell")
)
DEFAULT_LOCK = _RUNTIME_DIR / "zshell-battery-daemon.lock"


def _acquire_singleton_lock(path: Path) -> int | None:
	path.parent.mkdir(parents=True, exist_ok=True)
	fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o644)
	try:
		fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
	except BlockingIOError:
		os.close(fd)
		return None
	os.ftruncate(fd, 0)
	os.write(fd, str(os.getpid()).encode())
	return fd


@dataclass
class DeviceBattery:
	model: str
	nativePath: str | None
	serial: str | None
	type: str | None
	wired: bool
	percentage: float | None
	state: str | None
	charging: bool
	isPresent: bool


_SYSFS_POWER_SUPPLY = Path("/sys/class/power_supply")


def _normalize_serial(s: str) -> str:
	return re.sub(r"[^0-9A-Za-z]", "", s).upper()


def _find_upower_native_path(serial: str | None) -> str | None:
	if not serial or not _SYSFS_POWER_SUPPLY.is_dir():
		return None
	target = _normalize_serial(serial)
	for entry in _SYSFS_POWER_SUPPLY.glob("hidpp_battery_*"):
		try:
			raw = (entry / "serial_number").read_text().strip()
		except OSError:
			logger.debug(
				"couldn't read serial_number under %s", entry, exc_info=True
			)
			continue
		if _normalize_serial(raw) == target:
			return entry.name
	return None


def _process_battery_notification(dev, n) -> bool:

	if not getattr(dev, "isDevice", False):
		return False

	if int(n.sub_id) & 0x80:
		return False
	if n.sub_id == Notification.NO_OPERATION:
		return False

	if getattr(dev, "protocol", None) is not None and dev.protocol < 2.0:
		if n.sub_id in (Registers.BATTERY_STATUS, Registers.BATTERY_CHARGE):
			if n.data[-1:] != b"\x00":
				return False
			data = bytes([n.address]) + n.data
			dev.set_battery_info(hidpp10.parse_battery_status(n.sub_id, data))
			return True
		return False

	features = getattr(dev, "features", None)
	if not features:
		return False

	try:
		feature = features.get_feature(n.sub_id)
	except Exception:
		return False

	if feature == SupportedFeature.BATTERY_STATUS:
		if n.address == 0x00:
			dev.set_battery_info(hidpp20.decipher_battery_status(n.data)[1])
			return True
		return False

	if feature == SupportedFeature.BATTERY_VOLTAGE:
		if n.address == 0x00:
			dev.set_battery_info(hidpp20.decipher_battery_voltage(n.data)[1])
			return True
		return False

	if feature == SupportedFeature.UNIFIED_BATTERY:
		if n.address == 0x00:
			dev.set_battery_info(hidpp20.decipher_battery_unified(n.data)[1])
			return True
		return False

	if feature == SupportedFeature.ADC_MEASUREMENT:
		if n.address == 0x00:
			result = hidpp20.decipher_adc_measurement(n.data)
			if result:
				dev.set_battery_info(result[1])
				return True
		return False

	if feature == SupportedFeature.CENTURION_BATTERY_SOC:
		dev.set_battery_info(hidpp20.decipher_battery_centurion(n.data)[1])
		return True

	if feature == SupportedFeature.SOLAR_DASHBOARD:
		if n.data[5:9] == b"GOOD":
			charge, lux, _ = struct.unpack("!BHH", n.data[:5])
			status = BatteryStatus.DISCHARGING
			if n.address == 0x10 and lux > 200:
				status = BatteryStatus.RECHARGING
			dev.set_battery_info(Battery(charge, None, status, None, lux))
			return True
		return False

	return False


def _import_listener_deps():
	try:
		from logitech_receiver import base, device, receiver
		from logitech_receiver.listener import EventsListener
	except ImportError as e:
		raise RuntimeError(
			"logitech_receiver isn't importable -- install it with `pip install solaar`."
		) from e
	except ValueError as e:
		raise RuntimeError(
			"logitech_receiver failed to import because a GTK3 typelib is missing "
			f"({e}). Install the same GTK3 + PyGObject packages Solaar's GUI needs "
			"(e.g. python3-gi + gir1.2-gtk-3.0 on Debian/Ubuntu, or the gtk3/"
			"python3-gobject equivalents on your distro)."
		) from e

	class _BatteryEventsListener(EventsListener):
		def __init__(self, receiver_or_device, on_change):
			super().__init__(receiver_or_device, self._handle_notification)
			self._on_change = on_change

		def has_started(self):

			if not self.receiver.isDevice:
				try:
					self.receiver.notification_flags = (
						self.receiver.enable_connection_notifications()
					)
					self.receiver.notify_devices()
				except Exception:
					logger.exception(
						"failed enabling notifications for %s", self.receiver
					)

		def _handle_notification(self, n):
			if self.receiver.isDevice:
				try:
					_process_battery_notification(self.receiver, n)
				except Exception:
					logger.exception("processing %s for %s", n, self.receiver)
				self._on_change(self.receiver)
				return

			if n.devnumber == 0xFF:
				try:
					_process_battery_notification(self.receiver, n)
				except Exception:
					logger.exception("processing receiver notification %s", n)
				self._on_change(self.receiver)
				return

			if not (0 < n.devnumber <= 16):
				logger.warning(
					"unexpected device number %s in %s", n.devnumber, n
				)
				return

			try:
				dev = self.receiver[n.devnumber]
			except Exception:
				logger.exception(
					"resolving device number %s on %s",
					n.devnumber,
					self.receiver,
				)
				return

			if not dev:
				logger.warning(
					"%s: received %s for invalid device %d",
					self.receiver,
					n,
					n.devnumber,
				)
				return

			try:
				_process_battery_notification(dev, n)
			except Exception:
				logger.exception("processing %s for %s", n, dev)

			self._on_change(dev)

	return base, device, receiver, _BatteryEventsListener


def _start_hotplug_watcher(on_hotplug: Callable[[], None]):
	try:
		import pyudev
	except ImportError:
		logger.warning(
			"pyudev isn't importable -- hotplugged devices will only be "
			"noticed on the next --rescan-interval tick, not instantly. "
			"Install it with `pip install pyudev` (it's a core Solaar "
			"dependency, so it's normally already present)."
		)
		return None

	debounce_lock = threading.Lock()
	debounce_state: dict[str, threading.Timer | None] = {"timer": None}
	DEBOUNCE_SECONDS = 0.5

	def _fire():
		with debounce_lock:
			debounce_state["timer"] = None
		on_hotplug()

	def _handle_event(_device):
		if _device.action not in ("add", "remove"):
			return
		with debounce_lock:
			existing = debounce_state["timer"]
			if existing is not None:
				existing.cancel()
			timer = threading.Timer(DEBOUNCE_SECONDS, _fire)
			timer.daemon = True
			debounce_state["timer"] = timer
			timer.start()

	context = pyudev.Context()
	monitor = pyudev.Monitor.from_netlink(context)
	monitor.filter_by(subsystem="hidraw")
	observer = pyudev.MonitorObserver(monitor, callback=_handle_event)
	observer.start()
	return observer


class _ListenerRegistry:
	def __init__(self, on_change):
		self._on_change = on_change
		self._entries = {}
		self._last_change: dict[int, float] = {}
		self._lock = threading.Lock()
		self._rescan_lock = threading.Lock()
		self._min_rescan_interval = 1.0
		self._last_rescan_time = 0.0

		self._hidpp_io_lock = threading.RLock()

	def _touch_and_notify(self, dev_or_receiver, *args, **kwargs):
		with self._lock:
			self._last_change[id(dev_or_receiver)] = time.time()
		self._on_change(dev_or_receiver, *args, **kwargs)

	def last_change(self, dev) -> float:
		with self._lock:
			return self._last_change.get(id(dev), 0.0)

	def rescan(self):
		now = time.monotonic()
		if now - self._last_rescan_time < self._min_rescan_interval:
			return
		if not self._rescan_lock.acquire(blocking=False):
			return

		try:
			self._last_rescan_time = now
			with self._hidpp_io_lock:
				base, device, receiver, ListenerClass = _import_listener_deps()

				with self._lock:
					dead = [
						p
						for p, (_obj, listener) in self._entries.items()
						if not listener._active
					]
					for p in dead:
						dying_obj, _dying_listener = self._entries[p]

						if dying_obj.isDevice:
							ident = getattr(
								dying_obj, "unitId", None
							) or getattr(dying_obj, "serial", None)
							info = getattr(dying_obj, "battery_info", None)
							if ident and info is not None:
								for (
									other_path,
									(other_obj, _other_listener),
								) in self._entries.items():
									if other_path == p or other_obj.isDevice:
										continue
									for child in other_obj:
										child_ident = getattr(
											child, "unitId", None
										) or getattr(child, "serial", None)
										if child_ident == ident:
											child.set_battery_info(info)
											break
						del self._entries[p]
					known_paths = set(self._entries.keys())

				with self._lock:
					id_to_entry = {}
					for path, (obj, listener) in self._entries.items():
						if obj.isDevice:
							ident = obj.unitId or obj.serial
							if ident:
								id_to_entry[ident] = (path, obj, listener)

				for dev_info in base.receivers_and_devices():
					if dev_info.path in known_paths:
						continue
					try:
						if dev_info.isDevice:
							obj = device.create_device(base, dev_info)
						else:
							obj = receiver.create_receiver(base, dev_info)
					except OSError as e:
						if e.errno == 13:
							logger.error(
								"permission denied opening %s -- check the Solaar udev rule "
								"(rules.d/42-logitech-unify-permissions.rules) is installed",
								dev_info.path,
							)
						else:
							logger.exception("failed opening %s", dev_info)
						continue
					except Exception:
						logger.exception("failed opening %s", dev_info)
						continue

					if obj is None:
						continue

					listener = ListenerClass(obj, self._touch_and_notify)
					to_join = None

					if obj.isDevice:
						ident = None
						try:
							if obj.protocol >= 2.0:
								obj.get_ids()
							ident = obj.unitId or obj.serial
						except Exception:
							pass

						if ident:
							with self._lock:
								existing = id_to_entry.get(ident)
								if existing:
									old_path, old_obj, old_listener = existing

									if old_path != dev_info.path:
										logger.info(
											"replacing device %s (old path %s) with new path %s",
											ident,
											old_path,
											dev_info.path,
										)

										if (
											getattr(obj, "battery_info", None)
											is None
											and getattr(
												old_obj, "battery_info", None
											)
											is not None
										):
											obj.set_battery_info(
												old_obj.battery_info
											)

										old_listener.stop()

										del self._entries[old_path]

										self._entries[dev_info.path] = (
											obj,
											listener,
										)
										self._last_change[id(obj)] = time.time()

										to_join = old_listener
										break
									else:
										pass
								else:
									if (
										getattr(obj, "battery_info", None)
										is None
									):
										for (
											_,
											(other_obj, _other_listener),
										) in self._entries.items():
											if other_obj.isDevice:
												continue
											for child in other_obj:
												child_ident = getattr(
													child, "unitId", None
												) or getattr(
													child, "serial", None
												)
												if child_ident == ident:
													obj.set_battery_info(
														child.battery_info
													)
													break
											if (
												getattr(
													obj,
													"battery_info",
													None,
												)
												is not None
											):
												break

					with self._lock:
						if dev_info.path not in self._entries:
							self._entries[dev_info.path] = (obj, listener)
							self._last_change[id(obj)] = time.time()

					listener.start()

					logger.info(
						"listening on %s (%s)",
						dev_info.path,
						"device" if dev_info.isDevice else "receiver",
					)

					if to_join is not None:
						try:
							to_join.join(timeout=1.0)
						except Exception:
							logger.exception("error joining replaced listener")
						del to_join

		finally:
			self._rescan_lock.release()

	def known_devices(self):
		with self._lock:
			entries = list(self._entries.values())

		result = []
		seen_idents = set()
		for obj, _listener in entries:
			if obj.isDevice:
				ident = obj.unitId or obj.serial
				if ident and ident in seen_idents:
					continue
				if ident:
					seen_idents.add(ident)
				result.append(obj)
			else:
				for child in obj:
					ident = child.unitId or child.serial
					if ident and ident in seen_idents:
						continue
					if ident:
						seen_idents.add(ident)
					result.append(child)
		return result

	def stop(self):
		with self._lock:
			entries = list(self._entries.values())
		for _obj, listener in entries:
			try:
				listener.stop()
			except Exception:
				logger.exception("error stopping listener")
		for _obj, listener in entries:
			listener.join(timeout=2.0)


def _snapshot(registry: _ListenerRegistry) -> list[DeviceBattery]:
	results: list[DeviceBattery] = []

	with registry._hidpp_io_lock:
		for dev in registry.known_devices():
			info = dev.battery_info
			if info is None:
				try:
					if dev.ping():
						info = dev.battery()
				except Exception:
					logger.debug(
						"initial battery() failed for %s", dev, exc_info=True
					)
					info = None
			entry = _to_device_battery(dev, info)
			if entry is not None:
				results.append(entry)
	return results


def _clean_json(path: Path) -> None:
	_write_json_atomic(path, [])


def _iter_open_devices():
	try:
		from logitech_receiver import base, device, receiver
	except ImportError as e:
		raise RuntimeError(
			"logitech_receiver isn't importable -- install it with `pip install solaar`."
		) from e
	except ValueError as e:
		raise RuntimeError(
			"logitech_receiver failed to import because a GTK3 typelib is missing "
			f"({e}). Install the same GTK3 + PyGObject packages Solaar's GUI needs "
			"(e.g. python3-gi + gir1.2-gtk-3.0 on Debian/Ubuntu, or the gtk3/"
			"python3-gobject equivalents on your distro)."
		) from e

	for dev_info in base.receivers_and_devices():
		try:
			if dev_info.isDevice:
				d = device.create_device(base, dev_info)  # pyright: ignore[reportArgumentType]
				if d is not None:
					yield d
			else:
				d = receiver.create_receiver(base, dev_info)  # pyright: ignore[reportArgumentType]
				if d is not None:
					yield from d
		except OSError as e:
			if e.errno == 13:
				logger.error(
					"permission denied opening %s -- check the Solaar udev rule "
					"(rules.d/42-logitech-unify-permissions.rules) is installed "
					"and you're in the right group",
					dev_info.path,
				)
			else:
				logger.exception("failed opening %s", dev_info)
			continue
		except Exception:
			logger.exception("failed opening %s", dev_info)
			continue


def _to_device_battery(dev, battery) -> DeviceBattery | None:
	if battery is None:
		return None

	percentage = battery.level / 100 if isinstance(battery.level, int) else None
	state = battery.status.name.lower() if battery.status is not None else None

	serial = (
		getattr(dev, "serial", None) or getattr(dev, "unitId", None) or None
	)

	return DeviceBattery(
		model=dev.name or dev.codename or "Unknown device",
		nativePath=_find_upower_native_path(serial),
		serial=serial,
		type=str(dev.kind) if dev.kind is not None else None,
		wired=dev.receiver is None,
		percentage=percentage,
		state=state,
		charging=battery.charging(),
		isPresent=battery.ok(),
	)


def _read_battery(dev) -> DeviceBattery | None:
	try:
		if not dev.ping():
			return None
	except Exception:
		logger.debug("ping failed for %s", dev, exc_info=True)
		return None

	try:
		battery = dev.battery()
	except Exception:
		logger.debug("battery() failed for %s", dev, exc_info=True)
		return None

	return _to_device_battery(dev, battery)


def poll_once() -> list[DeviceBattery]:
	results: list[DeviceBattery] = []
	seen_idents = set()
	for dev in _iter_open_devices():
		try:
			info = _read_battery(dev)
			if info is not None:
				entry = _to_device_battery(dev, info)
				if entry:
					ident = entry.serial or (getattr(dev, "unitId", None))
					if ident and ident in seen_idents:
						continue
					if ident:
						seen_idents.add(ident)
					results.append(entry)
		finally:
			with contextlib.suppress(Exception):
				dev.close()
	return results


def _write_json_atomic(path: Path, results: list[DeviceBattery]) -> None:
	path.parent.mkdir(parents=True, exist_ok=True)
	payload = {
		"updated": time.time(),
		"devices": [asdict(r) for r in results],
	}
	fd, tmp_path = tempfile.mkstemp(
		dir=path.parent, prefix=".battery-", suffix=".tmp"
	)
	try:
		with os.fdopen(fd, "w") as f:
			json.dump(payload, f, indent=2)
		os.replace(tmp_path, path)
	except Exception:
		Path(tmp_path).unlink(missing_ok=True)
		raise


@app.command()
def daemon(
	rescan_interval: float = typer.Option(
		60.0,
		"--rescan-interval",
		"-i",
		help=(
			"Seconds between fallback rescans. This is NOT a battery poll "
			"interval -- battery updates are event-driven and written as "
			"soon as a device reports a change. New/removed devices "
			"(e.g. plugging a keyboard in to charge) are normally noticed "
			"within milliseconds via a udev hotplug watcher; this interval "
			"is only a safety net in case a udev event is ever missed."
		),
	),
	out: Path = typer.Option(
		DEFAULT_OUTPUT,
		"--out",
		"-o",
		help="Where to write the JSON status file.",
	),
	lock_file: Path = typer.Option(
		DEFAULT_LOCK,
		"--lock-file",
		help="Path used to ensure only one daemon runs at a time.",
	),
	verbose: bool = typer.Option(
		False,
		"--verbose",
		"-v",
		help="Log every listener start and battery-change write to stderr.",
	),
):
	lock_fd = _acquire_singleton_lock(lock_file)
	if lock_fd is None:
		try:
			holder_pid = lock_file.read_text().strip()
		except OSError:
			holder_pid = "unknown"
		typer.echo(
			f"error: a battery daemon is already running (pid {holder_pid}, lock: {lock_file})",
			err=True,
		)
		raise typer.Exit(code=1)

	if verbose:
		logging.basicConfig(
			level=logging.INFO, format="%(asctime)s %(name)s: %(message)s"
		)

	write_lock = threading.Lock()
	last_payload = None

	pending_write_timer: list[threading.Timer | None] = [None]
	pending_write_lock = threading.Lock()
	WRITE_SETTLE_SECONDS = 0.3

	write_seq_lock = threading.Lock()
	write_seq = 0
	last_written_seq = 0

	def write_if_changed():
		nonlocal last_payload, write_seq, last_written_seq
		with write_seq_lock:
			write_seq += 1
			my_seq = write_seq

		results = _snapshot(registry)
		with write_lock:
			with write_seq_lock:
				if my_seq <= last_written_seq:
					return
				last_written_seq = my_seq
			payload = [asdict(r) for r in results]
			if payload != last_payload:
				_write_json_atomic(out, results)
				last_payload = payload
				logger.info(
					"battery status changed, wrote %d device(s)", len(results)
				)

	def write_if_changed_coalesced():
		with pending_write_lock:
			existing = pending_write_timer[0]
			if existing is not None:
				existing.cancel()
			timer = threading.Timer(WRITE_SETTLE_SECONDS, write_if_changed)
			timer.daemon = True
			pending_write_timer[0] = timer
			timer.start()

	def on_change(_device_or_receiver, alert=None, reason=None):
		write_if_changed_coalesced()

	registry = _ListenerRegistry(on_change)

	def rescan_and_write():
		try:
			registry.rescan()
		except Exception:
			logger.exception("rescan failed")
		write_if_changed_coalesced()

	try:
		registry.rescan()
	except RuntimeError as e:
		typer.echo(f"error: {e}", err=True)
		raise typer.Exit(code=1) from None

	write_if_changed()

	hotplug_observer = _start_hotplug_watcher(rescan_and_write)

	running = True

	def _stop(signum, frame):
		nonlocal running
		running = False

	signal.signal(signal.SIGINT, _stop)
	signal.signal(signal.SIGTERM, _stop)

	typer.echo(
		f"Listening for battery events, writing to {out} (lock: {lock_file})"
	)

	try:
		while running:
			remaining = rescan_interval
			while running and remaining > 0:
				step = min(0.2, remaining)
				time.sleep(step)
				remaining -= step
			if running:
				rescan_and_write()
	finally:
		if hotplug_observer is not None:
			with contextlib.suppress(Exception):
				hotplug_observer.stop()
		with contextlib.suppress(Exception):
			_clean_json(out)

	registry.stop()
	typer.echo("Stopped.")


if __name__ == "__main__":
	app()