#!/usr/bin/python3

import gi
gi.require_version("Gtk", "3.0")

import dbus
import logging
import os
import setproctitle
import subprocess
import sys
import time
import threading
from dbus.mainloop.glib import DBusGMainLoop
from gi.repository import Gio, GLib, Gtk, Gdk, GObject
from Xlib import display, protocol, X, Xatom, error

class EWMH:
    """This class provides the ability to get and set properties defined
    by the EWMH spec. It was blanty ripped out of pyewmh
      * https://github.com/parkouss/pyewmh
    """

    def __init__(self, _display=None, root = None):
        self.display = _display or display.Display()
        self.root = root or self.display.screen().root

    def getActiveWindow(self):
        """Get the current active (toplevel) window or None (property _NET_ACTIVE_WINDOW)

        :return: Window object or None"""
        active_window = self._getProperty('_NET_ACTIVE_WINDOW')
        if active_window == None:
            return None

        return self._createWindow(active_window[0])

    def _getProperty(self, _type, win=None):
        if not win:
            win = self.root
        atom = win.get_full_property(self.display.get_atom(_type), X.AnyPropertyType)
        if atom:
            return atom.value

    def _setProperty(self, _type, data, win=None, mask=None):
        """Send a ClientMessage event to the root window"""
        if not win:
            win = self.root
        if type(data) is str:
            dataSize = 8
        else:
            data = (data+[0]*(5-len(data)))[:5]
            dataSize = 32

        ev = protocol.event.ClientMessage(window=win, client_type=self.display.get_atom(_type), data=(dataSize, data))

        if not mask:
            mask = (X.SubstructureRedirectMask|X.SubstructureNotifyMask)
        self.root.send_event(ev, event_mask=mask)

    def _createWindow(self, wId):
        if not wId:
            return None
        return self.display.create_resource_object('window', wId)

def format_path(path):
    logging.debug('Path:%s', path)
    result = path.replace('>', '', 1)
    result = result.replace('Root > ', '')
    result = result.replace('Label Empty > ', '')
    result = result.replace('_', '')
    return result.replace('>', u'\u0020\u0020\u00BB\u0020\u0020')

def get_bool(schema, path, key):
    if path:
        settings = Gio.Settings.new_with_path(schema, path)
    else:
        settings = Gio.Settings.new(schema)
    return settings.get_boolean(key)

def get_string(schema, path, key):
    if path:
        settings = Gio.Settings.new_with_path(schema, path)
    else:
        settings = Gio.Settings.new(schema)
    return settings.get_string(key)

def get_list(schema, path, key):
    if path:
        settings = Gio.Settings.new_with_path(schema, path)
    else:
        settings = Gio.Settings.new(schema)
    return settings.get_strv(key)

def rgba_to_hex(color):
   """
   Return hexadecimal string for :class:`Gdk.RGBA` `color`.
   """
   return "#{0:02x}{1:02x}{2:02x}".format(
                                    int(color.red   * 255),
                                    int(color.green * 255),
                                    int(color.blue  * 255))

def get_color(style_context, preferred_color, fallback_color):
    color = rgba_to_hex(style_context.lookup_color(preferred_color)[1])
    if color == '#000000':
        color = rgba_to_hex(style_context.lookup_color(fallback_color)[1])
    return color

def get_menu(menuKeys):
    """
    Generate menu of available menu items.
    """
    if not menuKeys:
        return ''

    menu_string, *menu_items = menuKeys
    for menu_item in menu_items:
        menu_string += '\n' + menu_item

    # Get the currently active font.
    font_name = get_string('org.mate.interface', None, 'font-name')

    # Get some colors from the currently selected theme.
    window = Gtk.Window()
    style_context = window.get_style_context()

    bg_color = get_color(style_context, 'dark_bg_color', 'theme_bg_color')
    fg_color = get_color(style_context, 'dark_fg_color', 'theme_fg_color')
    borders = get_color(style_context, 'borders', 'border_color')

    logging.debug('bg_color: %s', str(bg_color))
    logging.debug('fg_color: %s', str(fg_color))
    logging.debug('borders: %s', str(borders))

    selected_bg_color = rgba_to_hex(style_context.lookup_color('theme_selected_bg_color')[1])
    selected_fg_color = rgba_to_hex(style_context.lookup_color('theme_selected_fg_color')[1])
    error_bg_color = rgba_to_hex(style_context.lookup_color('error_bg_color')[1])
    error_fg_color = rgba_to_hex(style_context.lookup_color('error_fg_color')[1])
    info_bg_color = rgba_to_hex(style_context.lookup_color('info_bg_color')[1])
    info_fg_color = rgba_to_hex(style_context.lookup_color('info_fg_color')[1])
    text_color = rgba_to_hex(style_context.lookup_color('theme_text_color')[1])

    shortcut = get_shortcut()

    menu_cmd = subprocess.Popen(['rofi', '-dmenu', '-i',
                                 '-location', '1',
                                 '-width', '100', '-p', 'HUD: ',
                                 '-lines', '10', '-font', font_name,
                                 '-separator-style', 'none',
                                 '-hide-scrollbar',
                                 '-click-to-exit',
                                 '-levenshtein-sort',
                                 '-line-padding', '2',
                                 '-kb-cancel', 'Escape,' + shortcut,
                                 '-monitor', '-4', # show on monitor with the currently focused window
                                 '-color-enabled',
                                 '-color-window', bg_color +", " + borders + ", " + borders,
                                 '-color-normal', bg_color +", " + fg_color + ", " + bg_color + ", " + selected_bg_color + ", " + selected_fg_color,
                                 '-color-active', bg_color +", " + fg_color + ", " + bg_color + ", " + info_bg_color + ", " + info_fg_color,
                                 '-color-urgent', bg_color +", " + fg_color + ", " + bg_color + ", " + error_bg_color + ", " + error_fg_color],
                                 stdout=subprocess.PIPE, stdin=subprocess.PIPE)
    menu_cmd.stdin.write(menu_string.encode('utf-8'))
    menu_result = menu_cmd.communicate()[0].decode('utf8').rstrip()
    menu_cmd.stdin.close()

    return menu_result

"""
  try_appmenu_interface
"""
def try_appmenu_interface(window_id):
    # --- Get Appmenu Registrar DBus interface
    session_bus = dbus.SessionBus()
    try:
        appmenu_registrar_object = session_bus.get_object('com.canonical.AppMenu.Registrar', '/com/canonical/AppMenu/Registrar')
        appmenu_registrar_object_iface = dbus.Interface(appmenu_registrar_object, 'com.canonical.AppMenu.Registrar')
    except dbus.exceptions.DBusException:
        logging.info('Unable to register with com.canonical.AppMenu.Registrar.')
        return

    # --- Get dbusmenu object path
    try:
        dbusmenu_bus, dbusmenu_object_path = appmenu_registrar_object_iface.GetMenuForWindow(window_id)
    except dbus.exceptions.DBusException:
        logging.info('Unable to get dbusmenu object path.')
        return

    # --- Access dbusmenu items
    try:
        dbusmenu_object = session_bus.get_object(dbusmenu_bus, dbusmenu_object_path)
        dbusmenu_object_iface = dbus.Interface(dbusmenu_object, 'com.canonical.dbusmenu')
    except ValueError:
        logging.info('Unable to access dbusmenu items.')
        return

    dbusmenu_root_item = dbusmenu_object_iface.GetLayout(0, 0, ["label", "children-display"])
    dbusmenu_item_dict = dict()

    #For excluding items which have no action
    blacklist = []

    """ expanse_all_menu_with_dbus """
    def expanse_all_menu_with_dbus(item, root, path):
        item_id = item[0]
        item_props = item[1]

        # expand if necessary
        if 'children-display' in item_props:
            dbusmenu_object_iface.AboutToShow(item_id)
            dbusmenu_object_iface.Event(item_id, "opened", "not used", dbus.UInt32(time.time())) #fix firefox
        try:
            item = dbusmenu_object_iface.GetLayout(item_id, 1, ["label", "children-display"])[1]
        except:
            return

        item_children = item[2]

        if 'label' in item_props:
            new_path = path + " > " + item_props['label']
        else:
            new_path = path

        if len(item_children) == 0:
            if new_path not in blacklist:
                dbusmenu_item_dict[format_path(new_path)] = item_id
        else:
            blacklist.append(new_path)
            for child in item_children:
                expanse_all_menu_with_dbus(child, False, new_path)

    expanse_all_menu_with_dbus(dbusmenu_root_item[1], True, "")
    menuKeys = sorted(dbusmenu_item_dict.keys())

    menu_result = get_menu(menuKeys)

    # --- Use dmenu result
    if menu_result in dbusmenu_item_dict:
        action = dbusmenu_item_dict[menu_result]
        logging.debug('AppMenu Action : %s', str(action))
        dbusmenu_object_iface.Event(action, 'clicked', 0, 0)

    # Firefox:
    # Send closed events to level 1 items to make sure nothing weird happens
    # Firefox will close the submenu items (luckily!)
    # VimFx extension wont work without this
    dbusmenu_level1_items = dbusmenu_object_iface.GetLayout(0, 1, ["label"])[1]
    for item in dbusmenu_level1_items[2]:
        item_id = item[0]
        dbusmenu_object_iface.Event(item_id, "closed", "not used", dbus.UInt32(time.time()))

"""
  try_gtk_interface
"""
def try_gtk_interface(gtk_bus_name, gtk_menu_object_path, gtk_actions_paths_list):
    session_bus = dbus.SessionBus()
    # --- Ask for menus over DBus --- Credit @1931186
    try:
        gtk_menu_object = session_bus.get_object(gtk_bus_name, gtk_menu_object_path)
        gtk_menu_menus_iface = dbus.Interface(gtk_menu_object, dbus_interface='org.gtk.Menus')
    except dbus.exceptions.DBusException:
        logging.info('Unable to connect with com.gtk.Menus.')
        return

    # Here's the deal: The idea is to reduce the number of calls to the proxy and keep it as low as possible
    # because the proxy is a potential bottleneck
    # This means we ignore GMenus standard building model and just iterate over all the information one Start() provides at once
    # Start() does these calls, returns the result and keeps track of all parents (the IDs used by org.gtk.Menus.Start()) we called
    # queue() adds a parent to a potential_new_layers list; we'll use this later to avoid starting() some layers twice
    # explore is for iterating over the information a Start() call provides

    gtk_menubar_action_dict = dict()
    gtk_menubar_action_target_dict = dict()

    usedLayers = []
    def Start(i):
        usedLayers.append(i)
        return gtk_menu_menus_iface.Start([i])

    # --- Construct menu list ---

    potential_new_layers = []
    def queue(potLayer, label, path):
        # collects potentially new layers to check them against usedLayers
        # potLayer: ID of potential layer, label: None if nondescript, path
        potential_new_layers.append([potLayer, label, path])

    def explore(parent, path):
        for node in parent:
            content = node[2]
            # node[0] = ID of parent
            # node[1] = ID of node under parent
            # node[2] = actuall content of a node; this is split up into several elements/ menu entries
            for element in content:
                # We distinguish between labeled entries and unlabeled ones
                # Unlabeled sections/ submenus get added under to parent ({parent: {content}}), labeled under a key in parent (parent: {label: {content}})
                if 'label' in element:
                    if ':section' in element or ':submenu' in element:
                        # If there's a section we don't care about the action
                        # There theoretically could be a section that is also a submenu, so we have to handel this via queue
                        # submenus are more important than sections
                        if ':submenu' in element:
                            queue(element[':submenu'][0], None, path + " > " + element['label'])
                            # We ignore whether or not a submenu points to a specific index, shouldn't matter because of the way the menu got exportet
                            # Worst that can happen are some duplicates
                            # Also we don't Start() directly which could mean we get nothing under this label but this shouldn't really happen because there shouldn't be two submenus
                            # that point to the same parent. Even if this happens it's not that big of a deal.
                        if ':section' in element:
                            if element[':section'][0] != node[0]:
                                queue(element['section'][0], element['label'], path)
                                # section points to other parent, we only want to add the elements if their parent isn't referenced anywhere else
                                # We do this because:
                                # a) It shouldn't happen anyways
                                # b) The worst that could happen is we fuck up the menu structure a bit and avoid double entries
                    elif 'action' in element:
                        # This is pretty straightforward:
                        menu_action = str(element['action']).split(".",1)[1]
                        action_path = format_path(path + " > " + element['label'])
                        gtk_menubar_action_dict[action_path] = menu_action
                        if 'target' in element:
                            gtk_menubar_action_target_dict[action_path] = element['target']
                else:
                    if ':submenu' in element or ':section' in element:
                        if ':section' in element:
                            if element[':section'][0] != node[0] and element['section'][0] not in usedLayers:
                                queue(element[':section'][0], None, path)
                                # We will only queue a nondescript section if it points to a (potentially) new parent
                        if ':submenu' in element:
                            queue(element[':submenu'][0], None, path)
                            # We queue the submenu under the parent without a label

    queue(0, None, "")
    # We queue the first parent, [0]
    # This means 0 gets added to potential_new_layers with a path of "" (it's the root node)

    while len(potential_new_layers) > 0:
        layer = potential_new_layers.pop()
        # usedLayers keeps track of all the parents Start() already called
        if layer[0] not in usedLayers:
            explore(Start(layer[0]), layer[2])

    gtk_menu_menus_iface.End(usedLayers)

    menuKeys = sorted(gtk_menubar_action_dict.keys())
    menu_result = get_menu(menuKeys)

    # --- Use menu result
    if menu_result in gtk_menubar_action_dict:
        action = gtk_menubar_action_dict[menu_result]
        target = []
        try:
            target = gtk_menubar_action_target_dict[menu_result]
            if (not isinstance(target, list)):
                target = [target]
        except:
            pass

        for action_path in gtk_actions_paths_list:
            try:
                action_object = session_bus.get_object(gtk_bus_name, action_path)
                action_iface = dbus.Interface(action_object, dbus_interface='org.gtk.Actions')
                not_use_platform_data = dict()
                not_use_platform_data["not used"] = "not used"
                logging.debug('GTK Action : %s', str(action))
                action_iface.Activate(action, target, not_use_platform_data)
            except Exception as e:
                logging.debug('action_path: %s', str(action_path))

def hud(widget, keystr, user_data):
    logging.debug("Handling %s", str(user_data))

    # Get Window properties and GTK MenuModel Bus name
    ewmh = EWMH()
    win = ewmh.getActiveWindow()
    window_id = hex(ewmh._getProperty('_NET_ACTIVE_WINDOW')[0])
    gtk_bus_name = ewmh._getProperty('_GTK_UNIQUE_BUS_NAME', win)
    gtk_menu_object_path = ewmh._getProperty('_GTK_MENUBAR_OBJECT_PATH', win)
    gtk_app_object_path = ewmh._getProperty('_GTK_APPLICATION_OBJECT_PATH', win)
    gtk_win_object_path = ewmh._getProperty('_GTK_WINDOW_OBJECT_PATH', win)
    gtk_unity_object_path = ewmh._getProperty('_UNITY_OBJECT_PATH', win)

    gtk_bus_name, gtk_menu_object_path, gtk_app_object_path, gtk_win_object_path, gtk_unity_object_path = \
        [i.decode('utf8') if isinstance(i, bytes) \
        else i for i in [gtk_bus_name, gtk_menu_object_path, gtk_app_object_path, gtk_win_object_path, gtk_unity_object_path]]

    logging.debug('Window id is : %s', window_id)
    logging.debug('_GTK_UNIQUE_BUS_NAME: %s', gtk_bus_name)
    logging.debug('_GTK_MENUBAR_OBJECT_PATH: %s', gtk_menu_object_path)
    logging.debug('_GTK_UNIQUE_BUS_NAME: %s', gtk_bus_name)
    logging.debug('_GTK_MENUBAR_OBJECT_PATH: %s', gtk_menu_object_path)
    logging.debug('_GTK_APPLICATION_OBJECT_PATH: %s', gtk_app_object_path)
    logging.debug('_GTK_WINDOW_OBJECT_PATH: %s', gtk_win_object_path)
    logging.debug('_UNITY_OBJECT_PATH: %s', gtk_unity_object_path)

    if (not gtk_bus_name) or (not gtk_menu_object_path):
        logging.debug('Trying AppMenu')
        try_appmenu_interface(int(window_id, 16))
    else:
        # Many apps do not respect menu action groups, such as
        # LibreOffice and gnome-mpv, so we have to include all action
        # groups. Many other apps have these properties point to the
        # same path, so we need to remove them.
        logging.debug('Trying GTK interface')
        gtk_actions_paths_list = list(set([gtk_win_object_path,
                                  gtk_menu_object_path,
                                  gtk_app_object_path,
                                  gtk_unity_object_path]))
        try_gtk_interface(gtk_bus_name, gtk_menu_object_path, gtk_actions_paths_list)

class GlobalKeyBinding(GObject.GObject, threading.Thread):
    __gsignals__ = {
        'activate': (GObject.SignalFlags.RUN_LAST, None, ()),
    }

    def __init__(self):
        GObject.GObject.__init__(self)
        threading.Thread.__init__(self)
        self.setDaemon(True)

        self.display = display.Display()
        self.screen = self.display.screen()
        self.window = self.screen.root
        self.keymap = Gdk.Keymap().get_default()
        self.ignored_masks = self.get_mask_combinations(X.LockMask | X.Mod2Mask | X.Mod5Mask)
        self.map_modifiers()

    def get_mask_combinations(self, mask):
        return [x for x in range(mask+1) if not (x & ~mask)]

    def map_modifiers(self):
        gdk_modifiers = (Gdk.ModifierType.CONTROL_MASK, Gdk.ModifierType.SHIFT_MASK, Gdk.ModifierType.MOD1_MASK,
                         Gdk.ModifierType.MOD2_MASK, Gdk.ModifierType.MOD3_MASK, Gdk.ModifierType.MOD4_MASK, Gdk.ModifierType.MOD5_MASK,
                         Gdk.ModifierType.SUPER_MASK, Gdk.ModifierType.HYPER_MASK)
        self.known_modifiers_mask = 0
        for modifier in gdk_modifiers:
            if "Mod" not in Gtk.accelerator_name(0, modifier) or "Mod4" in Gtk.accelerator_name(0, modifier):
                self.known_modifiers_mask |= modifier

    def idle(self):
        self.emit("activate")
        return False

    def activate(self):
        GLib.idle_add(self.run)

    def grab(self, shortcut):
        accelerator = shortcut.replace("<Super>", "<Mod4>")
        keyval, modifiers = Gtk.accelerator_parse(accelerator)
        if not accelerator or (not keyval and not modifiers):
            self.keycode = None
            self.modifiers = None
            return False

        self.keycode = self.keymap.get_entries_for_keyval(keyval).keys[0].keycode
        self.modifiers = int(modifiers)

        # Request to receive key press/release reports from other windows that may not be using modifiers
        catch = error.CatchError(error.BadWindow)
        if self.modifiers:
            self.window.change_attributes(onerror=catch, event_mask=X.KeyPressMask|X.KeyReleaseMask)
        else:
            self.window.change_attributes(onerror=catch, event_mask=X.NoEventMask)
        if catch.get_error():
            return False

        catch = error.CatchError(error.BadAccess)
        for ignored_mask in self.ignored_masks:
            mod = self.modifiers | ignored_mask
            result = self.window.grab_key(self.keycode, mod, True, X.GrabModeAsync, X.GrabModeAsync, onerror=catch)
        # We grab Alt+click so that we can forward it to the window manager and allow Alt+click bindings (window move, resize, etc.)
        result = self.window.grab_button(X.AnyButton, X.Mod1Mask, True, X.ButtonPressMask, X.GrabModeSync, X.GrabModeAsync, X.NONE, X.NONE)
        self.display.flush()
        if catch.get_error():
            return False
        return True

    # Get which window manager we're currently using (Marco, Compiz, Metacity, etc...)
    def get_wm(self):
        name = ''
        wm_check = self.display.get_atom('_NET_SUPPORTING_WM_CHECK')
        win_id = self.window.get_full_property(wm_check, X.AnyPropertyType)
        if win_id:
            w = self.display.create_resource_object("window", win_id.value[0])
            wm_name = self.display.get_atom('_NET_WM_NAME')
            prop = w.get_full_property(wm_name, X.AnyPropertyType)
            if prop:
                name = prop.value
        return name.lower()

    def run(self):
        self.running = True
        wait_for_release = False
        while self.running:
            event = self.display.next_event()
            if event.type == X.KeyPress and event.detail == self.keycode and not wait_for_release:
                modifiers = event.state & self.known_modifiers_mask
                if modifiers == self.modifiers:
                    wait_for_release = True
            elif event.type == X.KeyRelease and event.detail == self.keycode and wait_for_release:
                GLib.idle_add(self.idle)
                wait_for_release = False
            elif event.type == X.ButtonPress:
                self.display.allow_events(X.ReplayPointer, X.CurrentTime)
                # Compiz would rather not have the event sent to it and just read it from the replayed queue
                # TODO: Explore changing Marco to behave similarly, thus eliminating the following chunk of code.
                #       This is not trivial and may break many other things.
                wm = self.get_wm()
                if wm != "compiz":
                    self.display.ungrab_keyboard(X.CurrentTime)
                    self.display.ungrab_pointer(X.CurrentTime)
                    query_pointer = self.window.query_pointer()
                    self.display.send_event(query_pointer.child, event, X.ButtonPressMask, True)
                wait_for_release = False
            else:
                if not self.modifiers:
                    self.display.ungrab_keyboard(X.CurrentTime)
                    self.display.send_event(self.window, event, X.KeyPressMask | X.KeyReleaseMask, True)
                wait_for_release = False

    def stop(self):
        self.running = False
        self.ungrab()
        self.display.close()

    def ungrab(self):
        if self.keycode:
            self.window.ungrab_key(self.keycode, X.AnyModifier, self.window)

    def rebind(self, shortcut):
        self.ungrab()
        if shortcut != "":
            self.grab(shortcut)

def get_shortcut():
    shortcut = 'Alt_L'
    try:
        shortcut = get_string('org.mate.hud', None, 'shortcut')
    except:
        logging.error('org.mate.hud gsettings not found. Defaulting to %s.' % shortcut)
    return shortcut

if __name__ == "__main__":
    setproctitle.setproctitle('mate-hud')
    logging.basicConfig(level=logging.INFO)

    def change_shortcut(schema, key):
        shortcut = settings.get_string("shortcut")
        keybinder.rebind(shortcut)

    shortcut = get_shortcut()

    DBusGMainLoop(set_as_default=True)
    keybinder = GlobalKeyBinding()
    keybinder.grab(shortcut)
    keybinder.connect("activate", hud, shortcut, "keystring %s (user data)" % shortcut)
    keybinder.start()
    logging.info("Press %s to handle keybinding", shortcut)

    settings = Gio.Settings.new("org.mate.hud")
    settings.connect("changed::shortcut", change_shortcut)

    try:
        GLib.MainLoop().run()
    except KeyboardInterrupt:
        GLib.MainLoop().quit()
