#!/usr/bin/env python

#  awesome-appmenu is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  sbomgr is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

#  You should have received a copy of the GNU General Public License
#  along with sbomgr.  If not, see <http://www.gnu.org/licenses/>.

#  Copyright (C) 2016 Daniel Prosser

import os, sys
from copy import copy

# Try to import configuration file
try:
  sys.path.insert(0, os.environ["HOME"] + "/.config/awesome-appmenu")
  import menurc
  config = True
except:
  config = False

################################################################################
# Launcher: finds and stores information for a .desktop file
class Launcher():

  # Exec field flags
  field_codes = ["%f", "%F", "%u", "%U", "%d", "%D", "%n", "%N", "%i", "%c", \
                 "%k", "%v", "%m"]

  def __init__(self, launcherfile, use_icon=True, 
               icon_search_paths=["/usr/share/icons/hicolor"], verbose=False):
    self.Name = None
    self.FileName = None
    self.FilePath = None
    self.Exec = None
    self.IconGenericName = None
    self.Icon = None
    self.NoDisplay = False
    self.OnlyShowIn = None
    self.Terminal = False
    self.Categories = []
    self.verbose = verbose
   
    self.getInfo(launcherfile, use_icon)
    if (use_icon and self.IconGenericName and not self.NoDisplay): 
      self.findIcon(icon_search_paths)
    self.expandFieldCodes()

  # Reads info from launcher file
  def getInfo(self, launcherfile, use_icon):

    self.FileName = launcherfile.split("/")[-1].strip()
    self.FilePath = launcherfile
    f = open(launcherfile)
    for line in f:
      splitline = line.split("=")

      if splitline[0] == "Name":
        # If Name is being defined again, just go on (google chrome and 
        # Libreoffice do this)
        if self.Name:
          if self.verbose:
            print("Redefinition of Name in " + launcherfile + ".")
          break
        else:
          self.Name = splitline[1].strip()
          if self.verbose:
            print("Name key for " + launcherfile + ":" + self.Name + ".")

      elif splitline[0] == "Exec":
        # Sometimes Exec lines have '=' if they set environment variables
        self.Exec = splitline[1].strip()
        for i in range(2, len(splitline)):
          self.Exec += "=" + splitline[i].strip()
        if self.verbose:
          print("Exec key for " + launcherfile + ":" + self.Exec + ".")

      elif splitline[0] == "Icon":
        if use_icon:
          # Full path specified
          if splitline[1][0] == "/": self.Icon = splitline[1].strip()
          # Name with or without extension
          else:
            if ( (splitline[1].strip()[-4:] == ".png") or 
                 (splitline[1].strip()[-4:] == ".svg") or
                 (splitline[1].strip()[-4:] == ".jpg") ):
              self.IconGenericName = splitline[1].strip()[:-4]
            else:
              self.IconGenericName = splitline[1].strip()
          if self.verbose:
            if self.IconGenericName:
              print("Icon key for " + launcherfile + ":" +
                    self.IconGenericName + ".")
            else:
              print("Icon key for " + launcherfile + ":" +
                    self.Icon + ".")

      elif splitline[0] == "Categories":
        self.Categories = splitline[1].split(";")
        if self.verbose:
          print("Categories key for " + launcherfile + ":" + splitline[1] + ".")
        for i in range(len(self.Categories)):
          self.Categories[i] = self.Categories[i].strip()

      elif splitline[0] == "NoDisplay":
        if (splitline[1].strip() == "true"):
          self.NoDisplay = True
          if self.verbose:
            print(launcherfile + " has NoDisplay set: skipping.")

      elif splitline[0] == "OnlyShowIn":
        self.OnlyShowIn = splitline[1].strip()
        if self.verbose:
          print("OnlyShowIn key for " + launcherfile + ": " +
                self.OnlyShowIn + ".")

      elif splitline[0] == "Terminal":
        if splitline[1].strip() == "true":
          if self.verbose:
            print(launcherfile + " launches in a terminal.")
          self.Terminal = True

    f.close()

  # Expands field codes in Exec key. See the freedesktop spec:
  # https://specifications.freedesktop.org/desktop-entry-spec/latest/ar01s06.html
  def expandFieldCodes(self):
    if not self.Exec: return
    for code in self.field_codes:

      if self.Exec.find(code) != -1:
        if code == "%c":

          # Replace any variants of %c, quoted or not, with name in quotes
          if self.Name: replacename = '"' + self.Name + '"'
          else: replacename = ""

          if self.Exec.find('"%c"') != -1: 
            self.Exec = self.Exec.replace('"%c"', replacename)
          if self.Exec.find("'%c'") != -1: 
            self.Exec = self.Exec.replace("'%c'", replacename)
          if self.Exec.find(code) != -1: 
            self.Exec = self.Exec.replace(code, replacename)

        elif code == "%k": 
          if self.FilePath: self.Exec = self.Exec.replace(code, self.FilePath)
          else: self.Exec = self.Exec.replace(code, "")
        elif code == "%i": 
          if (self.Icon):
            self.Exec = self.Exec.replace(code, "--icon " + self.Icon)
          else: self.Exec = self.Exec.replace(code, "")
   
        # Other field codes are not relevant for launching from appmenu
        else: self.Exec = self.Exec.replace(code, "")

    # Finally, remove extra spaces that may have been caused by removals
    splitexec = self.Exec.split()
    self.Exec = splitexec[0]
    for i in range(1, len(splitexec)):
      self.Exec += " " + splitexec[i] 

    if self.verbose:
      print("Expanded field code for " + self.Name + ": " + self.Exec + ".")

  # Finds icon by generic name
  def findIcon(self, icon_search_paths):

    # Note: this just takes the first matching icon that is encountered. It
    # doesn't prefer any particular size if there are multiple matches.
    for directory in icon_search_paths:

      # Skip any directory that doesn't exist
      if not os.path.isdir(directory):
        if self.verbose:
          print("Icon search path " + directory + " does not exist: skipping.")
        continue

      # Descend into directory
      for root, dirs, files in os.walk(directory):
        for name in files:
          if name[:-4] == self.IconGenericName:
            self.Icon = os.path.join(root, name)
            if self.verbose:
              print("Found icon for " + self.Name + ": " + self.Icon)
            break
        if self.Icon: break
      if self.Icon: break

################################################################################
# Category: groups and sorts launchers
class Category():

  def __init__(self, listed_name, name, use_icon=True, icon_generic_name=None,
               icon_search_paths=["/usr/share/icons/hicolor"], verbose=False):
    self.ListedName = listed_name
    self.Name = name
    self.UseIcon = use_icon
    self.IconGenericName = icon_generic_name
    self.Icon = None
    self.NumLaunchers = 0
    self.Launchers = []
    self.verbose = verbose

    # Determine if icon_generic_name supplied is actually an icon
    if use_icon and icon_generic_name:
      if ( (icon_generic_name.strip()[-4:] == ".png") or 
           (icon_generic_name.strip()[-4:] == ".svg") or
           (icon_generic_name.strip()[-4:] == ".jpg") ):
        self.Icon = icon_generic_name
        self.IconGenericName = None 
        if self.verbose:
          print("Using supplied path for category " + self.Name + " icon.")

    if (use_icon and self.IconGenericName): self.findIcon(icon_search_paths)

  # Finds icon by generic name
  def findIcon(self, icon_search_paths):

    # Note: this just takes the first matching icon that is encountered. It
    # doesn't prefer any particular size if there are multiple matches.
    for directory in icon_search_paths:

      # Skip any directory that doesn't exist
      if not os.path.isdir(directory):
        if self.verbose:
          print("Icon search path " + directory + " does not exist: skipping.")
        continue

      # Descend into directory
      for root, dirs, files in os.walk(directory):
        for name in files:
          if name[:-4] == self.IconGenericName:
            self.Icon = os.path.join(root, name)
            if self.verbose:
              print("Found icon for " + self.Name + " category: " + self.Icon)
            break
        if self.Icon: break
      if self.Icon: break

  # Adds a launcher
  def addLauncher(self, launcher):
    self.Launchers.append(launcher)
    self.NumLaunchers += 1
    if self.verbose:
      print("Added " + launcher.Name + " to " + self.Name + ".")

  # Sorts launchers
  def sortLaunchers(self):
    if self.verbose:
      print("Sorting launchers for " + self.Name + " category.")

    # Create a list of launchers by name, using FileName to distinguish
    launcherdict = {}
    for launcher in self.Launchers:
      launcherdict.update({launcher.FileName: launcher.Name})

    # Sort the list and create a new self.Launchers list
    # http://www.saltycrane.com/blog/2007/09/how-to-sort-python-dictionary-by-keys/
    templaunchers = copy(self.Launchers)
    self.Launchers = []
    for key, value in sorted(launcherdict.iteritems(), key=lambda(k,v): (v,k)):
      for launcher in templaunchers:
        if launcher.FileName == key: 
          self.Launchers.append(launcher)
          break

    if self.verbose:
      print("Finished sorting launchers for " + self.Name + " category.")

################################################################################
# Menu: groups categories
class Menu():

  def __init__(self, verbose=False):
    self.Categories = []
    self.NumCategories = 0
    self.miscidx = -1
    self.Favorites = []
    self.verbose = verbose

  # Adds a category
  def addCategory(self, category):
    self.Categories.append(category)
    self.NumCategories += 1
    if self.verbose: print("Added category " + category.Name)

  # Finds launchers and adds them to appropriate categories
  def getLaunchers(self, launcherpaths, use_icon=True, icon_search_paths=
                   ["/usr/share/icons/hicolor"], ignore_OnlyShowIn=False):

    # See if there's a Miscellaneous category or create it if needed
    for i in range(self.NumCategories):
      if self.Categories[i].Name == "Miscellaneous":
        self.miscidx = i
        break
    if self.miscidx == -1:
      self.addCategory(Category("Miscellaneous", "Miscellaneous", use_icon,
                                "applications-other", icon_search_paths,
                                verbose))
      self.miscidx = self.NumCategories - 1
      if self.verbose: print("Added Miscellaneous category idx {:d}.".format(
                              self.miscidx))

    # Find and group launchers into categories
    for directory in launcherpaths:

      # Skip any directory that doesn't exist
      if not os.path.isdir(directory): continue

      for root, dirs, files in os.walk(directory):
        for name in files:
          if name[-8:] == ".desktop":
            newlauncher = Launcher(os.path.join(root, name), use_icon,
                                   icon_search_paths, verbose)
            
            # Don't register launcher if there is no Exec field, if it has
            # NoDisplay=true, or if it has OnlyShowIn != awesome (unless
            # ignore_OnlyShowIn is enabled)
            if not newlauncher.Exec: 
              if self.verbose:
                print(newlauncher.Name + " has no Exec field: skipping.")
              continue
            if newlauncher.NoDisplay:
              if self.verbose:
                print(newlauncher.Name + " has NoDisplay set: skipping.")
              continue
            if not ignore_OnlyShowIn:
              if newlauncher.OnlyShowIn:
                if newlauncher.OnlyShowIn != "awesome":
                  if self.verbose:
                    print(newlauncher.Name + " is only shown in " +
                          newlauncher.OnlyShowIn + ": skipping.")
                  continue

            # Assign launcher to appropriate category(ies)
            self.assignLauncher(newlauncher)

  # Adds launcher to appropriate category(ies)
  def assignLauncher(self, launcher):

    launcher_grouped = False
    for category_name in launcher.Categories:
      for category in self.Categories:
        if category_name == category.ListedName:
          category.addLauncher(launcher)
          launcher_grouped = True
          if self.verbose:
            print("Assigned " + launcher.Name + " to " + category.Name +
                  " category.")

    # Put it in Miscellaneous category if nothing else fits
    if not launcher_grouped:
      self.Categories[self.miscidx].addLauncher(launcher)
      if self.verbose:
        print("Assigned " + launcher.Name + " to Miscellaneous category.")

  # Sorts launchers in each category
  def sortLaunchers(self):

    for category in self.Categories: category.sortLaunchers()

  # Creates favorites list
  def createFavoritesList(self, favorites):

    for favorite in favorites:
      found = False
      for category in self.Categories:
        for launcher in category.Launchers:
          if launcher.FileName == favorite:
            found = True
            self.Favorites.append(launcher)
            if self.verbose:
              print("Adding " + launcher.Name + " to favorites.")
            break
        if found: break
      if not found:
        print("Warning: could not find a launcher file named " + favorite + ".") 

  # Writes menu for awesomeWM
  def write(self, terminal="xterm"):

    home = os.environ["HOME"]
    menufile = home + "/.config/awesome/appmenu.lua"
    try:
      f = open(menufile, 'w')
    except IOError:
      os.makedirs(home + "/.config/awesome")
      f = open(menufile, 'w')

    # Header
    f.write("local appmenu = {}\n\n")

    # Write categories as long as they are not empty
    for category in self.Categories:
      if category.NumLaunchers == 0:
        if self.verbose:
          print(category.Name + " category is empty: it will not be written.")
        continue
      else:
        if self.verbose:
          print("Writing category " + category.Name + ".")
      f.write("appmenu." + category.Name + " = {\n")
      for launcher in category.Launchers:
        if self.verbose:
          print("Writing " + launcher.Name + " to " + category.Name + ".")
        if launcher.Terminal:
          if self.verbose:
            print(launcher.Name + " launches in a terminal.")
          f.write("    { '" + launcher.Name + "', '" + terminal + " -e " + 
                  launcher.Exec + "'")
        else: f.write("    { '" + launcher.Name + "', '" + launcher.Exec + "'")
        if launcher.Icon: 
          f.write(", '" + launcher.Icon + "'")
        f.write(" },\n")
      f.write("}\n\n")

    # Write favorites list if not empty
    if len(self.Favorites) != 0:
      if self.verbose:
        print("Writing favorites list.")
      f.write("appmenu.Favorites = {\n")
      for favorite in self.Favorites:
        if self.verbose:
          print("Writing " + favorite.Name + " to favorites.")
        if favorite.Terminal:
          if self.verbose:
            print(favorite.Name + " launches in a terminal.")
          f.write("    { '" + favorite.Name + "', '" + terminal + " -e " + 
                  favorite.Exec + "'")
        else: f.write("    { '" + favorite.Name + "', '" + favorite.Exec + "'")
        if favorite.Icon: 
          f.write(", '" + favorite.Icon + "'")
        f.write(" },\n")
      f.write("}\n\n")
    else:
      if self.verbose: print("Favorites list is empty: it will not be written.")

    # Write menu
    f.write("appmenu.Appmenu = {\n")
    for category in self.Categories:
      if category.NumLaunchers == 0: continue
      f.write("    { '" + category.Name + "', appmenu." + category.Name)
      if category.Icon:
        f.write(", '" + category.Icon + "'")
      f.write(" },\n")
    f.write("}\n\n")

    f.write("return appmenu")

    # Close file and print notification
    f.close()
    print("Wrote " + menufile + ".")

################################################################################
# Reads configuration file
def load_config():
  global config

  # Set defaults
  home = os.environ["HOME"]
  launcherpaths = ["/usr/share/applications", 
                   home + "/.local/share/applications"]
  iconpaths = ["/usr/share/icons/hicolor"]
  categories = [ ["Utility",   "Accessories", "applications-utilities"],
               ["Development", "Development", "applications-development"],
               ["Education",   "Education",   "applications-science"],
               ["Game",        "Games",       "applications-games"],
               ["Graphics",    "Graphics",    "applications-graphics"],
               ["Network",     "Internet",    "applications-internet"],
               ["Office",      "Office",      "applications-office"],
               ["AudioVideo",  "MultiMedia",  "applications-multimedia"],
               ["Settings",    "Settings",    "applications-accessories"],
               ["System",      "System",      "applications-system"],
               ["Wine",        "Wine",        "wine"] ]
  favorites = []
  terminal = "xterm"
  ignore_OnlyShowIn = False

  # Read configuration
  if config:

    try: launcherpaths = menurc.launcherpaths
    except AttributeError:
      print("Warning: menurc.py does not have launcherpaths. Using default.")

    try: iconpaths = menurc.iconpaths
    except AttributeError:
      print("Warning: menurc.py does not have iconpaths. Using default.")

    try: categories = menurc.categories
    except AttributeError:
      print("Warning: menurc.py does not have categories. Using default.")

    try: favorites = menurc.favorites
    except AttributeError: pass

    try: terminal = menurc.terminal
    except AttributeError: pass

    try: ignore_OnlyShowIn = menurc.ignore_OnlyShowIn
    except AttributeError: pass

  # Warnings about any non-existent directories
  for directory in launcherpaths:
    if not os.path.isdir(directory):
      print("Warning: search path " + directory + " does not exist:")
      print("No launchers will be searched for there.")

  for directory in iconpaths:
    if not os.path.isdir(directory):
      print("Warning: search path " + directory + " does not exist:")
      print("No icons will be searched for there.")

  return launcherpaths, iconpaths, categories, favorites, terminal, \
         ignore_OnlyShowIn

################################################################################
# Prints usage info
def print_help():

  print("Usage: " + sys.argv[0] + " OPTION")
  print("Options:")
  print("  --no-icons, -n: generate a menu without icons")
  print("  --verbose, -v:  print verbose output for debugging")
  print("  --help, -h:     show this usage information")

################################################################################
# Parses command line arguments
def parse_clos(argv):

  use_icon = True
  verbose = False
  for i in range(1, len(argv)):
    arg = argv[i]
    if (arg == "--help") or (arg == "-h"):
      print_help()
      exit(0)
    elif (arg == "--no-icons") or (arg == "-n"): use_icon = False
    elif (arg == "--verbose") or (arg == "-v"): verbose = True
    else:
      print("Unrecognized option " + arg + ".")
      print_help()
      exit(1)

  return use_icon, verbose

################################################################################
# Main program: generates an application menu for Awesome
if __name__ == "__main__":

  # Get command line arguments
  use_icon, verbose = parse_clos(sys.argv)

  # Load configuration file
  launcherpaths, iconpaths, categories, favorites, terminal, \
    ignore_OnlyShowIn = load_config()

  # Create a list of categories
  if use_icon: print("Finding category icons ...")
  appmenu = Menu(verbose)
  ncategories = len(categories)
  for i in range(ncategories):
    appmenu.addCategory(Category(categories[i][0], categories[i][1], use_icon,
                                 categories[i][2], iconpaths, verbose))

  # Find launchers and add them to appropriate categories
  if use_icon: print("Finding launcher files and icons ...")
  else:        print("Finding launchers ...")
  appmenu.getLaunchers(launcherpaths, use_icon, iconpaths, ignore_OnlyShowIn)

  # Sort launchers
  print("Sorting launchers ...")
  appmenu.sortLaunchers()

  # Get favorites menu
  if len(favorites) > 0:
    print("Creating favorites list ...")
    appmenu.createFavoritesList(favorites)

  # Write menu for awesomeWM
  appmenu.write(terminal=terminal)
