#!/usr/bin/env python

import os
import re
import socket
from libqtile.command.client import InteractiveCommandClient

layout_socket = '/tmp/qtile_kb_layout'
qtile_icc = InteractiveCommandClient()
current_keymap = 'us'
default_keymap = 'us'
windows_dict = {}

try:
    os.unlink(layout_socket)
except OSError:
    if os.path.exists(layout_socket):
        raise

# Create the Unix socket server
server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

# Bind the socket to the path
server.bind(layout_socket)

while True:
    server.listen(1)
    conn, addr = server.accept()
    data = ''
    while True:
        data_chunk = conn.recv(4096)
        if data_chunk:
            data += data_chunk.decode('UTF-8')
        if data in ['switch', 'win_new', 'win_focus']:
            break
        if not data_chunk or re.match(r'win_rm\s\d+', data):
            break
    if data == 'switch':
        if current_keymap == 'us':
            qtile_icc.core.set_keymap('ara')
            current_keymap = 'ara'
        elif current_keymap == 'ara':
            qtile_icc.core.set_keymap('us')
            current_keymap = 'us'
        try:
            # try to get current window id and store current keymap
            # in the windows dictionary
            wid = qtile_icc.window.info()['id']
            windows_dict[wid] = current_keymap
        except:
            pass
    elif data == 'win_new':
        try:
            wid = qtile_icc.window.info()['id']
            windows_dict[wid] = default_keymap
        except:
            pass
    elif data.startswith('win_rm'):
        try:
            wid = re.match(r'win_rm\s(\d+)', data).group(1)
            del windows_dict[int(wid)]
        except:
            pass
    elif data == 'win_focus':
        try:
            wid = qtile_icc.window.info()['id']
            if wid in windows_dict:
                qtile_icc.core.set_keymap(windows_dict[wid])
                current_keymap = windows_dict[wid]
            else:
                windows_dict[wid] = current_keymap
        except:
            pass

def layout_sock_conn(msg):
    layout_socket = '/tmp/qtile_kb_layout'
    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        client.connect(layout_socket)
        client.sendall(msg.encode())
        client.close()
    except:
        pass

@hook.subscribe.client_killed
def update_layout_tracker(win):
    layout_sock_conn(f'win_rm {win.wid}')

@hook.subscribe.client_new
def save_layout(win):
    layout_sock_conn('win_new')

@hook.subscribe.client_focus
@hook.subscribe.client_mouse_enter
def restore_layout(win):
    layout_sock_conn('win_focus')

def switch_layout(*args):
    layout_sock_conn('switch')

keys = [...
    Key([mod], "space", lazy.function(switch_layout), desc="switch keyboard layout"),
...]