#!/usr/bin/env python3

import json
import subprocess
import socket
import struct
import sys


_struct_header = '=6sII'

RUN_COMMAND = 0
GET_WORKSPACES = 1
SUBSCRIBE = 4
GET_OUTPUTS=3
GET_TREE = 4


def get_i3_socketpath():
    socketpath = subprocess.check_output(["i3", "--get-socketpath"])
    return socketpath.decode().rstrip()


def format_i3_command(kind: int, cmd: str):
    b = cmd.encode()
    return b"i3-ipc" + struct.pack("=II", len(b), kind) + b


def _ipc_recv(sock):
    data = sock.recv(14)

    if len(data) == 0:
        raise EOFError('got EOF from ipc socket')

    _, msg_length, msg_type = _unpack_header(data)
    msg_size = 6 + msg_length
    while len(data) < msg_size:
        data += sock.recv(msg_length)
    payload = _unpack(data)
    return payload, msg_type


def _unpack(data):
    _, msg_length, _ = _unpack_header(data)
    size = struct.calcsize(_struct_header)
    msg_size = size + msg_length
    payload = data[size:msg_size]
    return payload.decode('utf-8', 'replace')


def _unpack_header(data: bytes):
    return struct.unpack(_struct_header, data[:struct.calcsize(_struct_header)])


def open_i3_socket():
    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    client.connect(get_i3_socketpath())
    return client


def send_command(client, kind, data):
    client.sendall(format_i3_command(kind, data))
    result, _ = _ipc_recv(client)
    return json.loads(result)


class TreeNode:
    def __init__(self, json):
        self.id = None
        self.rect = None
        self.focused = False
        self.name = ""
        for k, v in json.items():
            if k != "nodes":
                self.__setattr__(k, v)
        self.parent = None
        self.nodes = list(map(TreeNode, json["nodes"]))
        for n in self.nodes:
            n.parent = self

    def find(self, where=None):
        if where is None:
            where = lambda _: True
        if where(self):
            yield self
        for n in self.nodes:
            yield from n.find(where)

    def __repr__(self) -> str:
        return "TreeNode(%d, %s, focused=%s, rect=%s)" % (self.id, self.name, self.focused, str(self.rect))

    def find_parent(self, where=None):
        if self.parent is None:
            return None
        if where is None or where(self.parent):
            return self.parent
        return self.parent.find_parent(where)

    def leaves(self):
        if len(self.nodes) == 0:
            yield self
        else:
            for n in self.nodes:
                yield from n.leaves()


def find_biggest_window(tree: TreeNode):
    max_leaf = None
    max_area = 0
    max_x = -1
    for leaf in tree.leaves():
        rect = leaf.rect
        area = rect["width"] * rect["height"]
        if area > max_area or (area == max_area and max_x > rect["x"]):
            max_area = area
            max_leaf = leaf
            max_x = rect["x"]

    return max_leaf


def current_workspace(tree):
    focused = list(tree.find(where=lambda n: n.focused is True))
    if len(focused) == 0:
        print("not found")
        sys.exit(1)
    focused = focused[0]
    ws = focused.find_parent(lambda n: n.type == "workspace")
    return ws


def promote(client):
    json = send_command(client, GET_TREE, "")
    tree = TreeNode(json)
    ws = current_workspace(tree)
    master = find_biggest_window(ws)
    send_command(client, RUN_COMMAND, "swap container with con_id %s" % master.id)


def main():
    if len(sys.argv) < 2:
        print("needs an argument")
        return
    command = sys.argv[1]
    client = open_i3_socket()
    if command == "promote":
        promote(client)
    client.close()

if __name__ == "__main__":
    main()
