tile

#!/usr/bin/env python3

import os
import sys
import time
import json
import fcntl
import ctypes
import ctypes.util

GAP = int(float(os.environ.get("TILE_GAP", "0")))
CHORD_WAIT = float(os.environ.get("TILE_CHORD_WAIT", "0.08"))
CHORD_HOLD = float(os.environ.get("TILE_CHORD_HOLD", "0.45"))
GTK_SHADOW = os.environ.get("TILE_GTK_SHADOW", "0") == "1"
WRAP = os.environ.get("TILE_WRAP", "0") == "1"

RUNTIME = os.environ.get("XDG_RUNTIME_DIR") or os.path.expanduser("~/.cache")
LOCK = os.path.join(RUNTIME, "tile-wm.lock")
STATE = os.path.join(RUNTIME, "tile-wm.state")

USAGE = """usage: tile <command>

  slot 1-4 - tile into quadrant N; focus the Nth window on the taskbar instead while Tab is held down
  focus 1-4 - focus the Nth window on the taskbar
  quad 1-4 - tile into quadrant N (hold several to span their bounding box)
  dir <left|right|up|down>
  full
  center
  status
"""

# quadrant -> (column, row)
QUAD = {1: (0, 0), 2: (0, 1), 3: (1, 0), 4: (1, 1)}
MINIMIZE = {2, 3}

SLOTS = ("1", "2", "3", "4")

# Window types the panel's task list never shows.
SKIP_TYPES = (
	"_NET_WM_WINDOW_TYPE_DESKTOP",
	"_NET_WM_WINDOW_TYPE_DOCK",
	"_NET_WM_WINDOW_TYPE_TOOLBAR",
	"_NET_WM_WINDOW_TYPE_MENU",
	"_NET_WM_WINDOW_TYPE_SPLASH",
	"_NET_WM_WINDOW_TYPE_DROPDOWN_MENU",
	"_NET_WM_WINDOW_TYPE_POPUP_MENU",
	"_NET_WM_WINDOW_TYPE_TOOLTIP",
	"_NET_WM_WINDOW_TYPE_NOTIFICATION",
	"_NET_WM_WINDOW_TYPE_COMBO",
	"_NET_WM_WINDOW_TYPE_DND"
)

ALL_DESKTOPS = 0xFFFFFFFF

Display, Window, Atom = ctypes.c_void_p, ctypes.c_ulong, ctypes.c_ulong

x11 = ctypes.CDLL(ctypes.util.find_library("X11"))
x11.XOpenDisplay.restype, x11.XOpenDisplay.argtypes = Display, [ctypes.c_char_p]
x11.XDefaultRootWindow.restype, x11.XDefaultRootWindow.argtypes = Window, [Display]
x11.XInternAtom.restype = Atom
x11.XInternAtom.argtypes = [Display, ctypes.c_char_p, ctypes.c_int]
x11.XGetWindowProperty.argtypes = [
	Display, Window, Atom, ctypes.c_long, ctypes.c_long, ctypes.c_int, Atom,
	ctypes.POINTER(Atom),
	ctypes.POINTER(ctypes.c_int),
	ctypes.POINTER(ctypes.c_ulong),
	ctypes.POINTER(ctypes.c_ulong),
	ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte))
]
x11.XFree.argtypes = [ctypes.c_void_p]
x11.XFlush.argtypes = [Display]
x11.XSync.argtypes = [Display, ctypes.c_int]
x11.XQueryKeymap.argtypes = [Display, ctypes.c_char * 32]
x11.XStringToKeysym.restype = ctypes.c_ulong
x11.XStringToKeysym.argtypes = [ctypes.c_char_p]
x11.XKeysymToKeycode.restype = ctypes.c_ubyte
x11.XKeysymToKeycode.argtypes = [Display, ctypes.c_ulong]
x11.XGetGeometry.argtypes = [
	Display, Window,
	ctypes.POINTER(Window),
	ctypes.POINTER(ctypes.c_int),
	ctypes.POINTER(ctypes.c_int),
	ctypes.POINTER(ctypes.c_uint),
	ctypes.POINTER(ctypes.c_uint),
	ctypes.POINTER(ctypes.c_uint),
	ctypes.POINTER(ctypes.c_uint)
]
x11.XTranslateCoordinates.argtypes = [
	Display, Window, Window, ctypes.c_int, ctypes.c_int,
	ctypes.POINTER(ctypes.c_int),
	ctypes.POINTER(ctypes.c_int),
	ctypes.POINTER(Window)
]


class XClientMessageEvent(ctypes.Structure):
	_fields_ = [
		("type", ctypes.c_int),
		("serial", ctypes.c_ulong),
		("send_event", ctypes.c_int),
		("display", Display),
		("window", Window),
		("message_type", Atom),
		("format", ctypes.c_int),
		("data", ctypes.c_long * 5)
	]


class XEvent(ctypes.Union):
	_fields_ = [
		("type", ctypes.c_int),
		("xclient", XClientMessageEvent),
		("pad", ctypes.c_long * 24)
	]


x11.XSendEvent.argtypes = [Display, Window, ctypes.c_int, ctypes.c_long, ctypes.POINTER(XEvent)]

dpy = x11.XOpenDisplay(None)
if not dpy:
	sys.exit("tile: cannot open X display")
root = x11.XDefaultRootWindow(dpy)

_atoms = {}


def atom(name):
	if name not in _atoms:
		_atoms[name] = x11.XInternAtom(dpy, name.encode(), False)
	return _atoms[name]


def prop(win, name):
	at, fmt = Atom(), ctypes.c_int()
	n, rem = ctypes.c_ulong(), ctypes.c_ulong()
	data = ctypes.POINTER(ctypes.c_ubyte)()
	rc = x11.XGetWindowProperty(dpy, win, atom(name), 0, 1024, False, 0, ctypes.byref(at), ctypes.byref(fmt), ctypes.byref(n), ctypes.byref(rem), ctypes.byref(data))

	if rc != 0 or not data:
		return None

	out = None

	if fmt.value == 32:
		arr = ctypes.cast(data, ctypes.POINTER(ctypes.c_ulong))
		out = [arr[i] for i in range(n.value)]

	x11.XFree(data)

	return out


def send_message(win, msg, *data):
	ev = XEvent()
	ev.xclient.type = 33  # ClientMessage
	ev.xclient.serial = 0
	ev.xclient.send_event = True
	ev.xclient.display = dpy
	ev.xclient.window = win
	ev.xclient.message_type = atom(msg)
	ev.xclient.format = 32

	for i in range(5):
		ev.xclient.data[i] = data[i] if i < len(data) else 0

	mask = (1 << 19) | (1 << 20)  # SubstructureNotify | SubstructureRedirect
	x11.XSendEvent(dpy, root, False, mask, ctypes.byref(ev))
	x11.XFlush(dpy)


def geometry(win):
	"""Absolute (x, y, w, h) of the window's client area."""
	r = Window()
	xx, yy = ctypes.c_int(), ctypes.c_int()
	w, h = ctypes.c_uint(), ctypes.c_uint()
	bw, d = ctypes.c_uint(), ctypes.c_uint()
	x11.XGetGeometry(dpy, win, ctypes.byref(r), ctypes.byref(xx), ctypes.byref(yy), ctypes.byref(w), ctypes.byref(h), ctypes.byref(bw), ctypes.byref(d))
	ax, ay, child = ctypes.c_int(), ctypes.c_int(), Window()
	x11.XTranslateCoordinates(dpy, win, root, 0, 0, ctypes.byref(ax), ctypes.byref(ay), ctypes.byref(child))
	return ax.value, ay.value, w.value, h.value


def keys_held(codes):
	buf = (ctypes.c_char * 32)()
	x11.XQueryKeymap(dpy, buf)
	raw = bytes(bytearray(buf))
	return {c for c in codes if raw[c >> 3] & (1 << (c & 7))}


def keycode(name):
	return x11.XKeysymToKeycode(dpy, x11.XStringToKeysym(name.encode()))


class Rect:
	__slots__ = ("x", "y", "w", "h")

	def __init__(self, x, y, w, h):
		self.x, self.y, self.w, self.h = int(x), int(y), int(w), int(h)

	@property
	def x2(self):
		return self.x + self.w

	@property
	def y2(self):
		return self.y + self.h

	@property
	def cx(self):
		return self.x + self.w / 2.0

	@property
	def cy(self):
		return self.y + self.h / 2.0

	def __repr__(self):
		return "%dx%d+%d+%d" % (self.w, self.h, self.x, self.y)


def overlap(a1, a2, b1, b2):
	return max(0, min(a2, b2) - max(a1, b1))


def screen_size():
	r = Window()
	xx, yy = ctypes.c_int(), ctypes.c_int()
	w, h = ctypes.c_uint(), ctypes.c_uint()
	bw, d = ctypes.c_uint(), ctypes.c_uint()
	x11.XGetGeometry(dpy, root, ctypes.byref(r), ctypes.byref(xx), ctypes.byref(yy), ctypes.byref(w), ctypes.byref(h), ctypes.byref(bw), ctypes.byref(d))
	return w.value, h.value


def monitors():
	try:
		xin = ctypes.CDLL(ctypes.util.find_library("Xinerama"))

		class ScreenInfo(ctypes.Structure):
			_fields_ = [
				("screen_number", ctypes.c_int),
				("x_org", ctypes.c_short),
				("y_org", ctypes.c_short),
				("width", ctypes.c_short),
				("height", ctypes.c_short)
			]

		xin.XineramaQueryScreens.restype = ctypes.POINTER(ScreenInfo)
		xin.XineramaQueryScreens.argtypes = [Display, ctypes.POINTER(ctypes.c_int)]
		n = ctypes.c_int()
		p = xin.XineramaQueryScreens(dpy, ctypes.byref(n))

		if p and n.value:
			mons = [Rect(p[i].x_org, p[i].y_org, p[i].width, p[i].height) for i in range(n.value)]
			x11.XFree(p)
			return sorted(mons, key=lambda m: (m.x, m.y))
	except Exception:
		pass

	w, h = screen_size()

	return [Rect(0, 0, w, h)]


def work_areas():
	"""Per-monitor work areas:"""
	sw, sh = screen_size()
	mons = monitors()
	areas = [Rect(m.x, m.y, m.w, m.h) for m in mons]

	clients = prop(root, "_NET_CLIENT_LIST") or []
	for c in clients:
		s = prop(c, "_NET_WM_STRUT_PARTIAL")
		if s and len(s) >= 12:
			left, right, top, bottom = s[0], s[1], s[2], s[3]
			ly1, ly2, ry1, ry2 = s[4], s[5], s[6], s[7]
			tx1, tx2, bx1, bx2 = s[8], s[9], s[10], s[11]
		else:
			s = prop(c, "_NET_WM_STRUT")
			if not s or len(s) < 4:
				continue
			left, right, top, bottom = s[0], s[1], s[2], s[3]
			ly1, ly2, ry1, ry2 = 0, sh, 0, sh
			tx1, tx2, bx1, bx2 = 0, sw, 0, sw

		for a in areas:
			if left and overlap(a.y, a.y2, ly1, ly2 + 1):
				nx = max(a.x, left)
				a.w -= nx - a.x
				a.x = nx

			if right and overlap(a.y, a.y2, ry1, ry2 + 1):
				a.w = min(a.x2, sw - right) - a.x

			if top and overlap(a.x, a.x2, tx1, tx2 + 1):
				ny = max(a.y, top)
				a.h -= ny - a.y
				a.y = ny

			if bottom and overlap(a.x, a.x2, bx1, bx2 + 1):
				a.h = min(a.y2, sh - bottom) - a.y
	return areas


def monitor_of(rect, areas):
	"""Index of the work area holding most rect"""
	best, best_area = 0, -1

	for i, a in enumerate(areas):
		o = overlap(rect.x, rect.x2, a.x, a.x2) * overlap(rect.y, rect.y2, a.y, a.y2)
		if o > best_area:
			best, best_area = i, o

	if best_area > 0:
		return best

	return min(range(len(areas)), key=lambda i: (areas[i].cx - rect.cx) ** 2 + (areas[i].cy - rect.cy) ** 2)


# noinspection shadowing-names
def neighbour(idx, direction, areas):
	"""Index of the adjacent monitor in `direction`."""
	cur = areas[idx]
	horiz = direction in ("left", "right")
	near = direction in ("left", "up")
	pos = (lambda i: areas[i].cx) if horiz else (lambda i: areas[i].cy)

	aligned = []
	for i, a in enumerate(areas):
		if i == idx:
			continue
		if horiz:
			if overlap(a.y, a.y2, cur.y, cur.y2) <= 0:
				continue
		elif overlap(a.x, a.x2, cur.x, cur.x2) <= 0:
			continue
		aligned.append(i)

	ahead = [i for i in aligned if (pos(i) < pos(idx) if near else pos(i) > pos(idx))]

	if ahead:
		return min(ahead, key=lambda i: abs(pos(i) - pos(idx)))
	if WRAP and aligned:
		return max(aligned, key=pos) if near else min(aligned, key=pos)

	return None


def cell_rect(area, cols, rows):
	halves_x = [area.x, area.x + area.w // 2, area.x2]
	halves_y = [area.y, area.y + area.h // 2, area.y2]

	x, x2 = halves_x[cols[0]], halves_x[cols[1]]
	y, y2 = halves_y[rows[0]], halves_y[rows[1]]

	if GAP:
		x += GAP if cols[0] == 0 else GAP // 2
		x2 -= GAP if cols[1] == 2 else GAP - GAP // 2
		y += GAP if rows[0] == 0 else GAP // 2
		y2 -= GAP if rows[1] == 2 else GAP - GAP // 2

	return Rect(x, y, x2 - x, y2 - y)


def detect_cell(frame, area):
	"""Which grid cell the window currently occupies"""
	spans = []

	mx = area.x + area.w // 2
	my = area.y + area.h // 2

	for lo, hi, size in ((area.x, mx, mx - area.x), (mx, area.x2, area.x2 - mx)):
		spans.append(overlap(frame.x, frame.x2, lo, hi) / float(size or 1))

	cols = [i for i, f in enumerate(spans) if f >= 0.62]
	spans = []

	for lo, hi, size in ((area.y, my, my - area.y), (my, area.y2, area.y2 - my)):
		spans.append(overlap(frame.y, frame.y2, lo, hi) / float(size or 1))

	rows = [i for i, f in enumerate(spans) if f >= 0.62]

	if not cols or not rows:
		return None, None

	return (cols[0], cols[-1] + 1), (rows[0], rows[-1] + 1)


def active_window():
	forced = os.environ.get("TILE_WINDOW")

	if forced:
		return int(forced, 0)

	p = prop(root, "_NET_ACTIVE_WINDOW")

	if not p or not p[0]:
		return None

	win = p[0]
	types = prop(win, "_NET_WM_WINDOW_TYPE") or []

	for bad in ("_NET_WM_WINDOW_TYPE_DESKTOP", "_NET_WM_WINDOW_TYPE_DOCK"):
		if atom(bad) in types:
			return None

	return win


def frame_rect(win):
	x, y, w, h = geometry(win)
	fe = prop(win, "_NET_FRAME_EXTENTS") or [0, 0, 0, 0]
	l, r, t, b = fe[0], fe[1], fe[2], fe[3]
	return Rect(x - l, y - t, w + l + r, h + t + b), (l, r, t, b)


def unmaximize(win):
	st = prop(win, "_NET_WM_STATE") or []
	wanted = [
		atom("_NET_WM_STATE_MAXIMIZED_VERT"),
		atom("_NET_WM_STATE_MAXIMIZED_HORZ"),
		atom("_NET_WM_STATE_FULLSCREEN")
	]

	if any(a in st for a in wanted):
		send_message(win, "_NET_WM_STATE", 0, wanted[0], wanted[1], 2)
		send_message(win, "_NET_WM_STATE", 0, wanted[2], 0, 2)
		x11.XSync(dpy, False)
		time.sleep(0.02)


def minimize(win):
	send_message(win, "WM_CHANGE_STATE", 3)
	x11.XSync(dpy, False)


def place(win, target):
	unmaximize(win)
	_, (l, r, t, b) = frame_rect(win)
	x, y, w, h = target.x, target.y, target.w, target.h

	if GTK_SHADOW:
		g = prop(win, "_GTK_FRAME_EXTENTS")
		if g and len(g) >= 4 and any(g[:4]):
			x -= g[0]
			y -= g[2]
			w += g[0] + g[1]
			h += g[2] + g[3]

	# With NorthWest gravity x/y address the frame, but w/h are the client size.
	flags = 0 | (1 << 8) | (1 << 9) | (1 << 10) | (1 << 11) | (2 << 12)
	send_message(win, "_NET_MOVERESIZE_WINDOW", flags, x, y, max(1, w - l - r), max(1, h - t - b))
	x11.XSync(dpy, False)


def read_state():
	try:
		with open(STATE) as f:
			return json.load(f)
	except Exception:
		return {}


def write_state(st):
	tmp = STATE + ".tmp"
	with open(tmp, "w") as f:
		json.dump(st, f)
	os.replace(tmp, STATE)


def cmd_quad(n):
	win = active_window()
	if not win:
		return

	now = time.time()
	st = read_state()
	keys = {n}

	if st.get("win") == win and now - st.get("t", 0) < CHORD_HOLD:
		prev = set(st.get("keys", []))
		keys |= prev
		if keys == prev:  # key auto-repeat
			st["t"] = now
			write_state(st)
			return
	else:
		# Wait briefly for keys pressed at the same instant as this one.
		codes = {q: keycode(str(q)) for q in QUAD}
		rev = {v: k for k, v in codes.items() if v}
		deadline = now + CHORD_WAIT
		while time.time() < deadline:
			for c in keys_held(set(rev)):
				keys.add(rev[c])

			if len(keys) > 1:
				break

			time.sleep(0.01)

	if keys == MINIMIZE:
		minimize(win)
		write_state({"t": time.time(), "win": win, "keys": sorted(keys)})
		return

	areas = work_areas()
	frame, _ = frame_rect(win)
	idx = monitor_of(frame, areas)
	cs = sorted(QUAD[k][0] for k in keys)
	rs = sorted(QUAD[k][1] for k in keys)
	cols, rows = (cs[0], cs[-1] + 1), (rs[0], rs[-1] + 1)
	place(win, cell_rect(areas[idx], cols, rows))
	write_state({"t": time.time(), "win": win, "keys": sorted(keys)})


def taskbar_windows():
	"""Windows the panel lists, in the order it lists them."""

	clients = prop(root, "_NET_CLIENT_LIST") or []
	cur = prop(root, "_NET_CURRENT_DESKTOP")
	cur = cur[0] if cur else 0
	skip = {atom(t) for t in SKIP_TYPES}
	out = []

	for c in clients:
		if skip & set(prop(c, "_NET_WM_WINDOW_TYPE") or []):
			continue

		if atom("_NET_WM_STATE_SKIP_TASKBAR") in (prop(c, "_NET_WM_STATE") or []):
			continue

		d = prop(c, "_NET_WM_DESKTOP")
		if d and d[0] != cur and d[0] != ALL_DESKTOPS:
			continue

		out.append(c)

	return out


def cmd_focus(n):
	wins = taskbar_windows()

	if not 1 <= n <= len(wins):
		return

	send_message(wins[n - 1], "_NET_ACTIVE_WINDOW", 2, 0, 0)
	write_state({})


def cmd_slot(n):
	tab = keycode("Tab")

	if tab and keys_held({tab}):
		cmd_focus(n)
	else:
		cmd_quad(n)


def cmd_dir(direction):
	win = active_window()
	if not win:
		return

	areas = work_areas()
	frame, _ = frame_rect(win)
	idx = monitor_of(frame, areas)
	cols, rows = detect_cell(frame, areas[idx])
	horiz = direction in ("left", "right")
	near = direction in ("left", "up")
	span = cols if horiz else rows
	other = rows if horiz else cols

	if span is None or other is None:  # untiled: half in that direction
		span, other = ((0, 1) if near else (1, 2)), (0, 2)
	elif span == (0, 2):  # full axis: shrink to that half
		span = (0, 1) if near else (1, 2)
	elif span == ((1, 2) if near else (0, 1)):  # opposite half: cross over
		span = (0, 1) if near else (1, 2)
	else:  # already at that edge
		nb = neighbour(idx, direction, areas)
		if nb is not None:
			idx = nb
			span = (1, 2) if near else (0, 1)  # enter from the far side
		else:
			other = (0, 2)  # no monitor there: fill the other axis

	cols, rows = (span, other) if horiz else (other, span)
	place(win, cell_rect(areas[idx], cols, rows))
	write_state({})


def cmd_full():
	win = active_window()
	if win:
		areas = work_areas()
		frame, _ = frame_rect(win)
		place(win, cell_rect(areas[monitor_of(frame, areas)], (0, 2), (0, 2)))
		write_state({})


def cmd_center():
	win = active_window()
	if win:
		areas = work_areas()
		frame, _ = frame_rect(win)
		a = areas[monitor_of(frame, areas)]
		place(win, Rect(a.x + (a.w - frame.w) // 2, a.y + (a.h - frame.h) // 2, frame.w, frame.h))
		write_state({})


def cmd_status():
	areas = work_areas()
	print("screen      : %dx%d" % screen_size())

	for i, (m, a) in enumerate(zip(monitors(), areas)):
		print("monitor %d   : %s   work area %s" % (i, m, a))

	win = active_window()

	if not win:
		print("active      : none")
		return

	frame, fe = frame_rect(win)
	idx = monitor_of(frame, areas)

	print("active      : 0x%08x  frame %s  extents l%d r%d t%d b%d" % (win, frame, fe[0], fe[1], fe[2], fe[3]))
	print("on monitor  : %d   cell %s" % (idx, detect_cell(frame, areas[idx])))


def main():
	args = sys.argv[1:]
	if not args:
		sys.exit(USAGE)

	cmd = args[0]
	lock = open(LOCK, "w")
	fcntl.flock(lock, fcntl.LOCK_EX)

	if cmd == "quad" and len(args) > 1 and args[1] in SLOTS:
		cmd_quad(int(args[1]))
	elif cmd == "slot" and len(args) > 1 and args[1] in SLOTS:
		cmd_slot(int(args[1]))
	elif cmd == "focus" and len(args) > 1 and args[1] in SLOTS:
		cmd_focus(int(args[1]))
	elif cmd == "dir" and len(args) > 1 and args[1] in ("left", "right", "up", "down"):
		cmd_dir(args[1])
	elif cmd == "full":
		cmd_full()
	elif cmd == "center":
		cmd_center()
	elif cmd == "status":
		cmd_status()
	else:
		sys.exit(USAGE)


if __name__ == "__main__":
	main()