#!/usr/bin/env python3

"""
Many applications put a delimiter and their name as a suffix in the window
title. Try to make it a lowercase prefix instead, with another delimiter, after
decluttering by removing everything but the executable name. We use different
dashes to detect delimiters, and executables in `PATH` as a heuristic to detect
names. We also detect when window titles already conform to our preferred
format, probably because the application was manually configured. Note that it
still has false positives. Also, this is maybe better done by modifying the
actual X11 window name.

Uses only the Python standard library.

References:

-   <https://i3wm.org/docs/ipc.html>
-   <https://docs.python.org/3/library/struct.html#byte-order-size-and-alignment>
-   <https://docs.python.org/3/library/struct.html#format-characters>
-   <https://i3wm.org/docs/userguide.html#pango_markup>
-   <https://docs.gtk.org/Pango/pango_markup.html>
-   <https://docs.gtk.org/glib/func.markup_escape_text.html>

You can run these tests with `i3-rename-windows --test`:

>>> rename('Some video - VLC media player')
'vlc:Some video'

>>> rename('Some website — Mozilla Firefox')
'firefox:Some website'

>>> rename('Mozilla Firefox')
'firefox'

>>> rename('Some other - Some unkown')
'Some other - Some unkown'

>>> rename('vim:the best editor - vim of course.txt')
'vim:the best editor - vim of course.txt'
"""

## Rename

import os
import re

SEP  = ':'
SEPS = '-–—'
EXES = '|'.join(
    re.escape(file)
    for path in os.environ.get('PATH', os.defpath).split(os.pathsep)
    for file in os.listdir(path)
    if os.access(os.path.join(path, file), os.F_OK | os.X_OK)
)
REGEX = re.compile(
    fr'^'
    fr'(?!(?:{EXES}){SEP}[^ ])'
    fr'(?:(?P<name>.*) +[{SEPS}] +)?'
    fr'.*'
    fr'(?P<exe>\b(?i:{EXES})\b)'
    fr'.*'
    fr'$'
)

def rename(name):
    match = re.match(REGEX, name)
    if match:
        return SEP.join(filter(None, [
            match['exe'].lower(),
            match['name'],
        ]))
    return name

## Serialization

import json
import struct

MAGIC   = b'i3-ipc'
HDR_FMT = '=LL'
HDR_LEN = struct.calcsize(HDR_FMT)
ESCAPE  = [
    ['\\', '\\\\' ],
    ['"',  '\\"'  ],
    ['&',  '&amp;'],
    ['<',  '&lt;' ],
    ['>',  '&gt;' ],
    ['%',  '&#37;'],
]

def recv(sock):
    assert sock.recv(len(MAGIC)) == MAGIC
    payload_len, tag = struct.unpack(HDR_FMT, sock.recv(HDR_LEN))
    return tag, json.loads(sock.recv(payload_len))

def send(sock, tag, msg=b'', sync_reply=False):
    payload = msg if type(msg) is bytes else json.dumps(msg).encode()
    header  = struct.pack(HDR_FMT, len(payload), tag)
    sock.sendall(MAGIC + header + payload)
    if sync_reply:
        reply_tag, reply_msg = recv(sock)
        assert reply_tag == tag
        return reply_msg

def escape(string):
    for old, new in ESCAPE:
        string = string.replace(old, new)
    return string

## Messages

MSG_COMMAND   = 0
MSG_SUBSCRIBE = 2
MSG_GET_TREE  = 4
EVT_WINDOW    = 3 | 1 << 31
EVT_SHUTDOWN  = 6 | 1 << 31

def command(sock, command):
    send(sock, MSG_COMMAND, command.encode())

def watch_windows(sock, changes, handle):
    def recurse(container):
        window = container['window']
        nodes  = container['nodes']
        if window:
            handle(sock, {'id': window, **container['window_properties']})
        for node in nodes:
            recurse(node)
    # TODO: Race? Subscribe first, handle the `MSG_GET_TREE` response in the
    # `while True`. This is our only use of `sync_reply`, so we could get rid
    # of that if we want.
    recurse(send(sock, MSG_GET_TREE, sync_reply=True))
    send(sock, MSG_SUBSCRIBE, ['window', 'shutdown'])
    while True:
        tag, msg = recv(sock)
        # TODO: `match`/`case`, requires newer Python.
        if tag == MSG_SUBSCRIBE:
            assert msg['success']
        elif tag == MSG_COMMAND:
            assert all(r['success'] for r in msg)
        elif tag == EVT_WINDOW:
            if msg['change'] in changes:
                recurse(msg['container'])
        elif tag == EVT_SHUTDOWN:
            return msg['change'] == 'restart'

## Titles

def replace(sock, props):
    props['title'] = escape(rename(props['title']))
    command(sock, '[id={id}] title_format "{title}"'.format(**props))

def restore(sock):
    # TODO: `[all]`, requires newer `i3`.
    command(sock, '[title=".*"] title_format "%title"')

## Main

import subprocess
import socket

def run(*args):
    return subprocess.check_output(args, text=True).strip()

def main():
    while True:
        with socket.socket(socket.AF_UNIX) as sock:
            sock.connect(run('i3', '--get-socket'))
            try:
                if not watch_windows(sock, ['new', 'title'], replace):
                    break
            except KeyboardInterrupt:
                break
            finally:
                restore(sock)
                sock.shutdown(socket.SHUT_RDWR)

if __name__ == '__main__':
    import sys
    args = iter(sys.argv[1:])
    if next(args, None) == '--test':
        import doctest
        exit(doctest.testmod().failed != 0)
    main()
