#!/usr/bin/python3
# -*- coding: utf-8 -*-
import sys
import os
import subprocess
import argparse
import re
import time
import traceback
from collections import namedtuple

ENABLE_LOGGING = False

# In the face of multiple screens with wildly different resolutions, there are
# essentially two ways to approach the 'grid size'.
# One is to take each screen and divide it into the same number of parts; so
# you have a laptop screen, and it gets divided into a 4x4 gride, and you have
# a 4K monitor, and it gets divided into a 4x4 grid.  This gives you the same
# size grid on both screens.
# The other is to divide the laptop screen into a 2x2 grid, and the 4K monitor
# into a 4x4 grid; this gives you similar size grid _cells_ on both screens.
# On my laptop monitor, I find a 4x4 grid to be awkwardly small.  So I think
# aiming for similar sized grid cells is going to be the better approach.  But
# then again, I could see a 3x3 grid on the laptop being reasonable, and is a
# significant improvement in flexibility.
#
# So.  I think this boils down to 'we need a configuration file'... though I'm
# not seeing a nice clean way to configure that...
# screen geometry -> AxB grid
#
# We want windows to snap to the grid for their location, and for their size.
# And we don't want windows to straddle screen boundaries.
# So... we could generate all possible grid entities, then search them to find
# the best target.

# Eventually, I'll move this to a configuration file, likely YAML, but since
# it's still in flux, just keep it here.
# I might want to support doing a 3x3 grid when it's the laptop screen only,
# but a 2x2 grid when it's with the 4k monitor
CONFIG = {
    'screen-grids': {
        # For a given screen, what grid to chop it into
        (1920, 1080): (2, 2), # Full HD
        (3840, 2160): (4, 4), # 4K UHD
    }
}


if ENABLE_LOGGING:
    LOG_FILE = open(os.path.expanduser('~/quicktile.log'), 'a')
    def LOG(message):
        message = message.rstrip()
        LOG_FILE.write("%s: %s\n" % (time.asctime(), message))
        LOG_FILE.flush()
else:
    def LOG(message):
        pass


def get_active_window_id():
    """gives the window ID of the currently active window"""
    task = subprocess.Popen(['xdotool', 'getactivewindow'], stdout=subprocess.PIPE)
    stdout, stderr = task.communicate()
    return int(stdout)


class Point(namedtuple('Point', ('X', 'Y'))):
    def distance_squared(self, other):
        return (self.X-other.X)**2 + (self.Y-other.Y)**2

    def __repr__(self):
        return "P(%s,%s)" % (self.X, self.Y)

    def __add__(self, other):
        summed = [s+o for s, o in zip(self, other)]
        return Point(*summed)

    def __sub__(self, other):
        diff = [s-o for s, o in zip(self, other)]
        return Point(*diff)

    def __mul__(self, factor):
        mul = [s * factor for s in self]
        return Point(*mul)

    def __rmul__(self, factor):
        return self * factor


class Geometry(namedtuple('Geometry', ('X', 'Y', 'W', 'H'))):
    def distance_squared(self, other):
        return (self.X-other.X)**2 + (self.Y-other.Y)**2

    def location_difference_squared(self, other):
        return (self.X-other.X)**2 + (self.Y-other.Y)**2

    def size_difference_squared(self, other):
        return (self.W-other.W)**2 + (self.H-other.H)**2

    def location_size_difference_squared(self, other):
        return (self.location_difference_squared(other), self.size_difference_squared(other))

    def size_location_difference_squared(self, other):
        """size is more important than distance"""
        return (self.size_difference_squared(other), self.location_difference_squared(other))

    def difference_squared(self, other):
        """returns square of distance between centers plus square of difference in size
        """
        return self.center().distance_squared(other.center()) + self.size_difference_squared(other)

    def __repr__(self):
        return "G(%s,%s,%s,%s)" % (self.X, self.Y, self.W, self.H)

    def __add__(self, other):
        summed = [s+o for s, o in zip(self, other)]
        return Geometry(*summed)

    def __sub__(self, other):
        diff = [s-o for s, o in zip(self, other)]
        return Geometry(*diff)

    def __mul__(self, factor):
        mul = [s * factor for s in self]
        return Geometry(*mul)

    def __rmul__(self, factor):
        return self * factor

    def nw(self):
        return Point(self.X, self.Y)

    def se(self):
        return Point(self.X+self.W, self.Y+self.H)

    def center(self):
        return Point(self.X+self.W//2, self.Y+self.H//2)

    def left_center(self):
        return Point(self.X, self.Y + self.H // 2)

    def right_center(self):
        return Point(self.X + self.W, self.Y + self.H // 2)

    def top_center(self):
        return Point(self.X + self.W // 2, self.Y)

    def bottom_center(self):
        return Point(self.X + self.W // 2, self.Y + self.H)


class Window(object):
    def __init__(self, id, desktop, X, Y, W, H, client, title):
        self.id = int(id, 16)
        self.desktop = int(desktop, 10)
        self.X = int(X, 10)
        self.Y = int(Y, 10)
        self.W = int(W, 10)
        self.H = int(H, 10)
        self.client = client
        self.title = title

    def _geometry_frame_offset(self):
        """Returns a Geometry for adjusting for the window frame.
        """
        # Normal maximized:
        #_KDE_NET_WM_FRAME_STRUT(CARDINAL) = 0, 0, 24, 0
        #_NET_FRAME_EXTENTS(CARDINAL) = 0, 0, 24, 0
        # Not maximized:
        #_KDE_NET_WM_FRAME_STRUT(CARDINAL) = 4, 4, 28, 4
        #_NET_FRAME_EXTENTS(CARDINAL) = 4, 4, 28, 4
        # left border, right border, title bar, bottom border
        task = subprocess.Popen(['xprop', '-id', str(self.id), '_NET_FRAME_EXTENTS'], stdout=subprocess.PIPE)
        stdout, stderr = task.communicate()
        left_border, right_border, title_bar, bottom_border = [int(v, 10) for v in stdout.decode().split('=')[-1].strip().split(', ')]
        return Geometry(-left_border, -title_bar, left_border+right_border, title_bar+bottom_border)

    def set_geometry(self, geometry):
        """Place a window at the given geometry, adjusting for window manager offsets.
        """
        orig_geometry = self.geometry()
        if geometry == orig_geometry: # Avoid work if it's already where we want it
            LOG( "Geometry already at %s" % (geometry, ))
        else:
            LOG( "Setting geometry to %s" % (geometry, ))
            # NOTE: If the window is maximized, the xdotool will not be able to
            # move/resize the window, and will hang for 15 seconds.

            # We can detect a normal, maximized window:
            # _NET_WM_STATE(ATOM) =  _NET_WM_STATE_MAXIMIZED_VERT, _NET_WM_STATE_MAXIMIZED_HORZ
            # Windows can have just one of those set, so we detect either
            task = subprocess.Popen(['xprop', '-id', str(self.id), '_NET_WM_STATE'], stdout=subprocess.PIPE)
            stdout, stderr = task.communicate()
            flags = set(stdout.decode().split('=')[-1].strip().split(', '))
            LOG("flags=%r" % flags)
            if flags.intersection(('_NET_WM_STATE_MAXIMIZED_VERT', '_NET_WM_STATE_MAXIMIZED_HORZ')):
                LOG("DETECTED MAXIMIZED WINDOW")
                subprocess.check_call(['wmctrl', '-i', '-r', hex(self.id), '-b', 'remove,maximized_vert,maximized_horz'])
                time.sleep(0.10) # Give it a (longer) moment to resize
            # But KDE's "Quick Tile" feature does not set anything that xprop
            # displays, so it won't detect windows that have been "quick tiled".

            frame_offset = self._geometry_frame_offset()
            LOG("frame_offset=%s" % (frame_offset,))
            offset_geometry = geometry - frame_offset
            # offset_geometry does not include the frame
            if self.client != 'N/A': # Bug workaround
                # For windows that have the hostname as the client instead of
                # 'N/A', when we set the geometry we have to provide the NW
                # corner _including_ the frame, but the width and height
                # _without_ the frame.
                offset_geometry += Geometry(frame_offset.X, frame_offset.Y, 0, 0)
            LOG("frame_offset fixup=%s" % (frame_offset,))
            self._set_geometry(offset_geometry)
            new_geometry = self.get_geometry()
            if new_geometry == orig_geometry:
                # It didn't move.  One of the ways this can happen is when KDE's
                # native quick tiling or maximizing is in use on a window.
                LOG( "Geometry unchanged; attempting to unmaximize")
                self.unmaximize()
                self._set_geometry(offset_geometry)
                new_geometry = self.get_geometry()
            if new_geometry != geometry: # The window manager is being a real pain
                LOG( "\007Failed to set geometry to %s using %s; wound up with %s instead" % (geometry, offset_geometry, new_geometry))

    def _set_geometry(self, geometry):
        """Directly calls an xdotool command to size and move the window to the given coordinates.
        """
        # Move vs size order matters.  If shrinking, size then move.  If growing, move then size.
        if geometry.W > self.geometry().W or geometry.H > self.geometry().H: # Growing
            cmd = ['xdotool', 'windowmove', '--sync', str(self.id), str(geometry.X), str(geometry.Y), 'windowsize', '--sync', str(self.id), str(geometry.W), str(geometry.H)]
        else: # Shrinking
            cmd = ['xdotool', 'windowsize', '--sync', str(self.id), str(geometry.W), str(geometry.H), 'windowmove', '--sync', str(self.id), str(geometry.X), str(geometry.Y)]
        start = time.time()
        subprocess.check_call(cmd)
        end = time.time()
        if end-start > 1:
            LOG("\007_set_geometry took %0.2fs" % (end-start))
        time.sleep(0.05) # Without this sleep, xdotool will _sometimes_ fail to set the geometry.

    def geometry(self):
        return self._get_geometry() + self._geometry_frame_offset()

    def get_geometry(self):
        return self.geometry()

    def _get_geometry(self):
        task = subprocess.Popen(['xdotool', 'getwindowgeometry', str(self.id)], stdout=subprocess.PIPE)
        stdout, stderr = task.communicate()
        geoRE = re.compile('Window .*Position: (?P<X>-?[0-9]+),(?P<Y>-?[0-9]+) .*Geometry: (?P<W>[0-9]+)x(?P<H>[0-9]+)', re.DOTALL)
        match = geoRE.match(stdout.decode('utf-8'))
        result = dict((k, int(v)) for k, v in match.groupdict().items())
        return Geometry(**result)

    def unmaximize(self):
        # Removing maximization is not sufficient; we have to add it and then
        # remove it to get it to un-maximize.  Has the annoying side effect
        # when a window is 'KDE quicktiled' to a half or quarter of the screen
        # of maximizing the window for a brief flash before resizing to the
        # desired size.
        subprocess.check_call(['wmctrl', '-i', '-r', hex(self.id), '-b', 'add,maximized_vert,maximized_horz'])
        subprocess.check_call(['wmctrl', '-i', '-r', hex(self.id), '-b', 'remove,maximized_vert,maximized_horz'])
        time.sleep(0.05)

    def __repr__(self):
        return 'W(%s,%s,%s,%s,%s,%s,%s,%s)' % (self.id, self.desktop, self.X, self.Y, self.W, self.H, self.client, self.title)


def get_current_windows():
    task = subprocess.Popen(['wmctrl', '-l', '-G'], stdout=subprocess.PIPE)
    stdout, stderr = task.communicate()
    entry = re.compile('(?P<id>0x[0-9a-f]{8})\\s+(?P<desktop>[-0-9]+)\\s+'
        '(?P<X>-?[0-9]+)\\s+'
        '(?P<Y>-?[0-9]+)\\s+'
        '(?P<W>[0-9]+)\\s+'
        '(?P<H>[0-9]+)\\s+'
        '(?P<client>[^ ]+)\\s(?P<title>.*)', re.M)
    windows = []
    for line in stdout.splitlines():
        match = entry.match(line.decode('utf-8'))
        if not match:
            raise Exception("Failed to parse %r" % line)
        windows.append(Window(**match.groupdict()))
    return windows


class Controller(object):
    def __init__(self):
        """Where grid is the number of cells in each direction on each desktop.
        KDE's default quick tiling is equivalent to grid=2.
        """
        self.query_window_manager()
        self._calculate_geometries()

    def query_window_manager(self):
        self.current_windows = get_current_windows()
        self.windows_by_id = dict((w.id, w) for w in self.current_windows)
        self.active = get_active_window_id()

    def _generate_grid_cells_for_desktop(self, geometry, grid_x, grid_y):
        """Returns a set of Geometry objects for all possible grid placements
        on the given geometry, when divided into a grid_x-by-grid_y grid.
        """
        grid_cells = []
        for grid_left in range(grid_x):
            for grid_right in range(grid_left, grid_x):
                for grid_top in range(grid_y):
                    for grid_bottom in range(grid_top, grid_y):
                        left = geometry.X + geometry.W * grid_left // grid_x
                        right = geometry.X + geometry.W * (grid_right+1) // grid_x
                        top = geometry.Y + geometry.H * grid_top // grid_y
                        bottom = geometry.Y + geometry.H * (grid_bottom+1) // grid_y
                        width = right - left
                        height = bottom - top
                        grid = Geometry(X=left,
                                        Y=top,
                                        W=width,
                                        H=height)
                        grid_cells.append(grid)
        return grid_cells

    def _splits_for_screen_size(self, width, height):
        _, closest = min(((width-w)**2 + (height-h)**2, (w, h)) for w, h in CONFIG['screen-grids'].keys())
        splits = CONFIG['screen-grids'].get(closest)
        if not splits:
            splits = (2, 2)
            LOG("No split info found for %sx%s screen, defaulting to %sx%s" % (width, height, splits[0], splits[1]))
        return splits

    def _calculate_geometries(self):
        self.plasma_windows = [w for w in self.current_windows if w.desktop == -1]
        menubar = [w for w in self.plasma_windows if w.title == 'Plasma'][0]
        desktops = [w for w in self.plasma_windows if w.title.startswith('Desktop')]
        # KDE's logic for moving windows into a given location seems to think
        # that the menubar is on all desktops, and will refuse to move a window
        # down far enough to cover it, and will move the window if it is
        # resized enough to cover it.
        if False:
            # For now, I'm going to assume the menu bar is on the bottom of the screen.
            menubar_desktop = [w for w in desktops if w.X == menubar.X][0]
            other_desktops = [w for w in desktops if w != menubar_desktop]
            self.desktop_geometries = [
                Geometry(menubar_desktop.X, menubar_desktop.Y, menubar_desktop.W, menubar_desktop.H-menubar.H),
            ]
            self.desktop_geometries.extend(Geometry(w.X, w.Y, w.W, w.H) for w in other_desktops)
        else:
            # So we're going to simply give up on the bottom 28px of the non-menubar screen, and act like it exists on all screens
            self.desktop_geometries = [Geometry(d.X, d.Y, d.W, d.H-menubar.H) for d in desktops]
        #LOG("Desktop geometries: %s\n" % self.desktop_geometries) # DEBUG

        self.grid_tiles = []
        for desktop_geometry in self.desktop_geometries:
            x_splits, y_splits = self._splits_for_screen_size(desktop_geometry.W, desktop_geometry.H)
            self.grid_tiles.extend(self._generate_grid_cells_for_desktop(desktop_geometry, x_splits, y_splits))
        #LOG("GRID_TILES: %s" % self.grid_tiles) # DEBUG

    def active_window(self):
        return self.windows_by_id[self.active]

    def window_action(self, action, direction):
        window = self.active_window()
        original_window_geometry = window.geometry()
        LOG("Starting geometry = %s" % (original_window_geometry, ))
        _, window_geometry = min([(original_window_geometry.location_size_difference_squared(g), g) for g in self.grid_tiles])
        LOG("Snapped geometry = %s" % (window_geometry, ))
        # Need to figure out the window's nearest grid-granular size
        command = (action, direction)
        # Move
        # Only consider cells that overlap the area directly in the direction
        # of the desired motion.  This means that a window at the top of one
        # screen, when pushed up, won't move horizontally to another screen
        # that is 'higher'.
        if command == ('move', 'left'):
            grids = [(window_geometry.right_center().distance_squared(g.right_center()) + window_geometry.size_difference_squared(g), g)
                for g in self.grid_tiles if g.nw().X < window_geometry.nw().X and g.se().X < window_geometry.se().X
                    and g.nw().Y < window_geometry.se().Y and g.se().Y > window_geometry.nw().Y]
        elif command == ('move', 'right'):
            grids = [(window_geometry.left_center().distance_squared(g.left_center()) + window_geometry.size_difference_squared(g), g)
                for g in self.grid_tiles if g.nw().X > window_geometry.nw().X and g.se().X > window_geometry.se().X
                    and g.nw().Y < window_geometry.se().Y and g.se().Y > window_geometry.nw().Y]
        elif command == ('move', 'up'):
            grids = [(window_geometry.bottom_center().distance_squared(g.bottom_center()) + window_geometry.size_difference_squared(g), g)
                for g in self.grid_tiles if g.nw().Y < window_geometry.nw().Y and g.se().Y < window_geometry.se().Y
                    and g.nw().X < window_geometry.se().X and g.se().X > window_geometry.nw().X]
        elif command == ('move', 'down'):
            grids = [(window_geometry.top_center().distance_squared(g.top_center()) + window_geometry.size_difference_squared(g), g)
                for g in self.grid_tiles if g.nw().Y > window_geometry.nw().Y and g.se().Y > window_geometry.se().Y
                    and g.nw().X < window_geometry.se().X and g.se().X > window_geometry.nw().X]
        # Grow
        elif command == ('grow', 'left'):
            grids = [(window_geometry.X - g.X, g) for g in self.grid_tiles if g.H == window_geometry.H and g.W > window_geometry.W and g.se() == window_geometry.se()]
        elif command == ('grow', 'right'):
            grids = [(g.se().X - window_geometry.se().X, g) for g in self.grid_tiles if g.H == window_geometry.H and g.W > window_geometry.W and g.nw() == window_geometry.nw()]
        elif command == ('grow', 'up'):
            grids = [(window_geometry.Y - g.Y, g) for g in self.grid_tiles if g.H > window_geometry.H and g.W == window_geometry.W and g.se() == window_geometry.se()]
        elif command == ('grow', 'down'):
            grids = [(g.se().Y - window_geometry.se().Y, g) for g in self.grid_tiles if g.H > window_geometry.H and g.W == window_geometry.W and g.nw() == window_geometry.nw()]
        # Shrink
        elif command == ('shrink', 'left'):
            grids = [(window_geometry.se().X - g.se().X, g) for g in self.grid_tiles if g.H == window_geometry.H and g.W < window_geometry.W and g.nw() == window_geometry.nw()]
        elif command == ('shrink', 'right'):
            grids = [(g.X - window_geometry.X, g) for g in self.grid_tiles if g.H == window_geometry.H and g.W < window_geometry.W and g.se() == window_geometry.se()]
        elif command == ('shrink', 'up'):
            grids = [(window_geometry.se().Y - g.se().Y, g) for g in self.grid_tiles if g.H < window_geometry.H and g.W == window_geometry.W and g.nw() == window_geometry.nw()]
        elif command == ('shrink', 'down'):
            grids = [(g.Y - window_geometry.Y, g) for g in self.grid_tiles if g.H < window_geometry.H and g.W == window_geometry.W and g.se() == window_geometry.se()]
        # Snap
        elif command == ('snap', 'here'):
            grids = [(window_geometry.location_size_difference_squared(g), g) for g in self.grid_tiles]
        else:
            raise Exception("Bad command %s %s" % (action, direction))

        if not grids:
            LOG("No target identified, finding closest tile.")
            grids = [(window_geometry.difference_squared(g), g) for g in self.grid_tiles]
        else:
            LOG("Sorted qualified grids: %s" % sorted(grids))
        _difference, grid = min(grids)
        window.set_geometry(grid)


def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=['move', 'grow', 'shrink', 'snap'])
    parser.add_argument('direction', choices=['left', 'right', 'up', 'down', 'here'])
    args = parser.parse_args(argv[1:])

    LOG("start %s %s" % (args.action, args.direction))
    try:
        controller = Controller()
        controller.window_action(args.action, args.direction)
    except Exception as error:
        LOG("ERROR: %s" % error)
        LOG(traceback.format_exc())
        raise
    LOG("end %s %s" % (args.action, args.direction))

    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))