#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#
# Author: Serg Kolo , contact: 1047481448@qq.com
# Date: September 27 , 2016
# Purpose: appindicator for displaying mounted filesystem usage
# Tested on: Ubuntu 16.04 LTS
#
#
# Licensed under The MIT License (MIT).
# See included LICENSE file or the notice below.
#
# Copyright © 2016 Sergiy Kolodyazhnyy
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import gi
gi.require_version('AyatanaAppIndicator3', '0.1')
gi.require_version('Notify', '0.7')
from gi.repository import GLib as glib
from gi.repository import AyatanaAppIndicator3 as appindicator
from gi.repository import Gtk as gtk
from gi.repository import Gio
from gi.repository import Notify
from os import statvfs
#from collections import OrderedDict
import subprocess
import copy
import shutil
import dbus
import math
import json
import os

class UdisksIndicator(object):

    def __init__(self):
        self.app = appindicator.Indicator.new(
            'udisks-indicator', "drive-harddisk-symbolic.svg",
            appindicator.IndicatorCategory.HARDWARE
            )

        if not self.app.get_icon():
           self.app.set_icon("drive-harddisk-symbolic")

        self.app.set_status(appindicator.IndicatorStatus.ACTIVE)


        self.fields = ['Partition','Alias','Drive',
                  'MountPoint','Filesystem','Usage'
        ]
        
        self.user_home = os.path.expanduser('~')
        filename = '.udisks-indicator-preferences.json'
        self.prefs_file = os.path.join(self.user_home,filename)
        self.prefs = self.read_prefs_file() 
        #self.fields_set = self.prefs['fields']
        #self.autostart = self.prefs[]

        filename = '.partition_aliases.json'
        self.config_file = os.path.join(self.user_home,filename)
        self.cache = self.get_partitions()
        
        self.note = Notify.Notification.new(__file__,None,None)
        Notify.init(__file__)

        self.cache_icon_theme = self.gsettings_get(
             'org.gnome.desktop.interface',None,'icon-theme' 
        )
        self.make_menu()
        self.update()

    def read_prefs_file(self,*args):
        default = { 'fields': [True]*5, 'autostart':False }
        if not os.path.exists(self.prefs_file):
           return default

        with open(self.prefs_file) as f:
            try:
                return json.load(f)
            except:
                return default
 
    def write_prefs_file(self,*args):
        with open(self.prefs_file,'w') as f:
            try:
                json.dump(self.prefs,f,
                          indent=4,sort_keys=True
                )
            except Exception as e:
                self.send_notif(
                    self.note,
                    'Failed writing '+self.prefs_file,
                    str(e)
                )
                
                
 

    def update(self):
        timeout = 5
        glib.timeout_add_seconds(timeout,self.callback)

    def callback(self):
        current = self.get_partitions()
        icon_theme = self.gsettings_get(
                     'org.gnome.desktop.interface',None,'icon-theme' 
        )
        if self.cache != current:
            self.cache = current
            self.make_menu()
        if self.cache_icon_theme != icon_theme:
            self.cache_icon_theme = icon_theme
            self.make_menu()
        self.update()        

    def add_menu_item(self,menu_obj,item_type,image,label,action,args):
        """ dynamic function that can add menu items depending on
            the item type and other arguments"""
        menu_item,icon = None,None
        if item_type is gtk.ImageMenuItem and label:
            menu_item = gtk.ImageMenuItem.new_with_label(label)
            menu_item.set_always_show_image(True)
            if '/' in image:
                icon = gtk.Image.new_from_file(image)
            else:
                icon = gtk.Image.new_from_icon_name(image,48)
            menu_item.set_image(icon)
        elif item_type is gtk.ImageMenuItem and not label:
            menu_item = gtk.ImageMenuItem()
            menu_item.set_always_show_image(True)
            if '/' in image:
                icon = gtk.Image.new_from_file(image)
            else:
                icon = gtk.Image.new_from_icon_name(image,16)
            menu_item.set_image(icon)
        elif item_type is gtk.MenuItem:
            menu_item = gtk.MenuItem(label)
        elif item_type is gtk.SeparatorMenuItem:
            menu_item = gtk.SeparatorMenuItem()
        if action:
            menu_item.connect('activate',action,*args)

        menu_obj.append(menu_item)
        menu_item.show()

    def get_fields(self,*args):
        """ Returns info field items necessary for
            each entry, padded depending on 
            the system font used - default or monospace.
            Other font types are not covered
        """
        fields = [ i + ':' for i in self.fields]
        
        
        if args[-1] == 'default':
            # Thanks to Mateo for figuring out these numbers
            fields[0] = fields[0] + " "*11
            fields[1] = fields[1] + " "*20
            fields[2] = fields[2] + " "*19
            fields[3] = fields[3] + " "*4
            fields[4] = fields[4] + " "*8
            fields[5] = fields[5] + " "*18

        if args[-1] == 'mono':
            for index,item in enumerate(fields):
                fields[index] = self.pad_string(item,20," ")

        return fields


    def make_menu(self,*args):
        """ generates entries in the indicator"""
        if hasattr(self, 'app_menu'):
            for item in self.app_menu.get_children():
                self.app_menu.remove(item)
        self.app_menu = gtk.Menu()
        
        # Determine font first so we know
        # how to pad the text
        font = self.gsettings_get(
                    'org.gnome.desktop.interface',
                    None,
                    'font-name'
        )
        font = str(font).replace("'","") 

        fields = []
        if 'Ubuntu' in font and font.split()[1].isdigit():
            fields = self.get_fields('default')
        elif 'Mono' in font:
            fields = self.get_fields('mono')
        else:
            fields = self.get_fields('neither')


        drive_icon = 'gnome-dev-harddisk'
        usb_icon = 'media-removable'
        optical_icon = 'media-optical'

        icon_theme = self.gsettings_get(
                     'org.gnome.desktop.interface',None,'icon-theme' 
        )
        if str(icon_theme) == "'ubuntukylin-icon-theme'":
           usb_icon = 'drive-harddisk-usb'

         
        # MOUNTED 
        # list of tuples,(partition,drive,mountpoint,usage)
        partitions = self.get_partitions()
        for i in partitions:
            label_lines = [None]*6
            label_lines[0] = fields[0] + i[0]
            label_lines[2] = fields[2] + i[1]
            label_lines[3] = fields[3] + i[2]
            label_lines[5] = fields[5] + i[3] + '%'
            
            hr_value = self.get_human_readable(int(i[4]))
            label_lines[5] = label_lines[5] + ' (' + hr_value + ')' 

            icon = drive_icon
            #icon = 'gnome-dev-harddisk' 
            #icon = 'block-device'
            #icon = 'drive-harddisk'
            #icon = 'drive-harddisk-symbolic.svg'
            if self.get_bus_type(i[1]) == 'usb':
                icon = usb_icon
               #icon = 'drive-harddisk-usb'
               #icons = 'gnome-dev-removable-usb'
               #icon = 'media-removable'
               #icon = 'drive-harddisk-usb'
               #icon = 'drive-removable-media'
               #icon = 'gnome-dev-removable'
               #icon = 'media-flash'

            fs_type = str(self.get_filesystem_type(i[0]))
            if fs_type:
               if 'iso' in fs_type:
                   icon = optical_icon
               label_lines[4] = fields[4] +  fs_type
            else:
               label_lines[4] = fields[4] +  "unknown"

            alias = self.find_alias(i[0])
            if alias:
                label_lines[1] = fields[1] + alias
            else:
                label_lines[1] = None

            usage_bar = self.make_usage_bar(i[3])
            label_lines.append(usage_bar)

            for j,f in enumerate(self.prefs['fields']):
                if not f:
                   label_lines[j] = None
            

            label = '\n'.join([l for l in label_lines if l])

            contents = [self.app_menu,gtk.ImageMenuItem,icon,
                        label,self.open_mountpoint,[i[2]]
            ]
            self.add_menu_item(*contents)       

            contents = [self.app_menu,gtk.SeparatorMenuItem,None,
                        None,None,[None]
            ]
            self.add_menu_item(*contents)

        self.unmounted = gtk.MenuItem('Unmounted Partitions')
        self.unmounted_submenu = gtk.Menu()
        self.unmounted.set_submenu(self.unmounted_submenu)

        # UNMOUNTED
        for i in self.get_unmounted_partitions():

            # TODO: add type checking, prevent swap
            # org.freedesktop.UDisks2.Block.IdType

            part = "Partition: " + i[0]
            alias = self.find_alias(i[0])
            drive = "\nDrive: " + i[1]
            fs_type = str(self.get_filesystem_type(i[0]))

            icon = drive_icon
            if fs_type:
               filesystem = "\nFilesystem: " + fs_type
               #icon = 'drive-removable-media'
               #icon = 'drive-multidisk'
               #icon = 'drive-harddisk-system'
               if fs_type == 'swap':
                  filesystem = filesystem + " NOT MOUNTABLE"
                  icon = 'dialog-error'
               if 'iso' in fs_type:
                  icon = optical_icon
            else:
               filesystem = "\nFilesystem: unknown"
            
            if  self.get_bus_type(i[1]) == 'usb':
                icon = 'drive-removable-media-usb'

            label = part + drive + filesystem

            if alias: 
               alias = "\nAlias: " + alias
               label = part + alias + drive + filesystem

            contents = [
                 self.unmounted_submenu,gtk.ImageMenuItem,icon,
                 label,self.mount_partition,[i[0]]
            ]
            self.add_menu_item(*contents)
            contents = [
                 self.unmounted_submenu,gtk.SeparatorMenuItem,None,
                 None,None,[None]
            ]
            self.add_menu_item(*contents)

        self.app_menu.append(self.unmounted)
        self.unmounted.show()

        
        self.separator = gtk.SeparatorMenuItem()
        self.app_menu.append(self.separator)
        self.separator.show()

        contents = [self.app_menu,gtk.ImageMenuItem,'text-editor',
                    'Make Alias',self.make_alias,[None]
        ]
        self.add_menu_item(*contents)
 
#        desktop_file = '.config/autostart/udisks-indicator.desktop'
#        full_path = os.path.join(self.user_home,desktop_file)
#        label = 'Start Automatically' 

        #menu_item = gtk.CheckMenuItem.new_with_label(label)
        #if os.path.exists(full_path):
        #     menu_item.set_active(True)   
        #menu_item.connect('activate',self.toggle_auto_startup)
        #self.app_menu.append(menu_item)
        #menu_item.show()

        #   label = ' \u2714' + label
        #contents = [self.app_menu,gtk.CheckMenuItem,None,
        #            label,self.toggle_auto_startup,[None]
        #]
        #self.add_menu_item(*contents)
     

        contents = [self.app_menu,gtk.ImageMenuItem,'gnome-disks',
                    'Open Disks Utility',self.open_disks_utility,[None]
        ] 
        self.add_menu_item(*contents) 
        contents = [self.app_menu,gtk.ImageMenuItem,'gtk-preferences',
                    'Preferences',self.preferences_dialog,[None]
        ]
        self.add_menu_item(*contents)
        contents = [self.app_menu,gtk.ImageMenuItem,'exit',
                    'Quit',self.quit,[None]
        ]
        self.add_menu_item(*contents)
        self.app.set_menu(self.app_menu)

    def preferences_dialog(self,*args):

        def _toggle_field(*args):
            pref_vals = args[-1]
            pref_vals[args[-2]] = args[0].get_state()
            #self.make_menu()

        def _ok_action(*args):
            self.prefs['fields'] = copy.deepcopy(args[-2])
            self.prefs['autostart'] = args[-1][0]
            self.process_autostart_file()
            self.write_prefs_file()
            glib.timeout_add_seconds(1,self.make_menu)
            window.hide()

        def _set_defaults(*args):
            self.prefs['fields'] = [True]*5
            self.prefs['autostart'] = False
            os.unlink(self.prefs_file)
            window.hide()

        def _toggle_autostart(*args):
            obj = args[-1]
            obj[0] = args[0].get_state()

        def _cancel_action(*args):
            window.hide()

       
        prefs = copy.deepcopy(self.prefs['fields'])
        autostart = [copy.deepcopy(self.prefs['autostart'])]
        window = gtk.Window()
        window.set_border_width(40)
        vert = gtk.Orientation.VERTICAL
        horz = gtk.Orientation.HORIZONTAL

        main_box = gtk.Box(orientation=vert,spacing=20)
        listbox = gtk.ListBox()
        listbox.set_selection_mode(gtk.SelectionMode.NONE)
        main_box.pack_start(listbox,True,True,0)
        window.add(main_box)

        row = gtk.ListBoxRow()
        listbox.add(row)
        label = gtk.Label('',xalign=0)
        label.set_markup("<b>INFO FIELDS</b>")
        row.add(label)

        for i,f in enumerate(self.fields): 
            if i == 0 or i == 5:
                continue
            row = gtk.ListBoxRow()
            listbox.add(row)
            aux_box1 = gtk.Box(orientation=horz, spacing=20)
            aux_box2 = gtk.Box(orientation=vert,spacing=20)
            aux_box1.pack_start(aux_box2,True,True,0)
           
            switch = gtk.Switch()
            label = gtk.Label(f,xalign=0)

            switch.props.valign = gtk.Align.CENTER
            if self.prefs['fields'][i]:
                switch.set_state(True)
            switch.connect('notify::active',_toggle_field,i,prefs)
    
            aux_box1.pack_start(switch,False,True,0)
            aux_box2.pack_start(label,True,True,0)
            row.add(aux_box1)

        row = gtk.ListBoxRow()
        listbox.add(row)
        label = gtk.Label('',xalign=0)
        label.set_markup("<b>OTHER OPTIONS</b>")
        row.add(label)

 
        row = gtk.ListBoxRow()
        listbox.add(row)
        aux_box = gtk.Box(orientation=horz,spacing=20)
        switch = gtk.Switch()
        if self.prefs['autostart']:
           switch.set_state(True)
        switch.connect('notify::active',_toggle_autostart,autostart)
        label = gtk.Label('Autostart',xalign=0)
        aux_box.pack_start(switch,False,True,0)
        aux_box.pack_start(label,True,True,0)
        row.add(aux_box)

        row = gtk.ListBoxRow()
        listbox.add(row)
        aux_box = gtk.Box(orientation=horz,spacing=20)
        for i in ['Set Defaults','Cancel','OK']:
            button = gtk.Button.new_with_label(i)
            if i == 'OK':
               button.connect('clicked',_ok_action,prefs,autostart)
            if i == 'Cancel':
               button.connect('clicked',_cancel_action,autostart)
            if i == 'Set Defaults':
               button.connect('clicked',_set_defaults)
            aux_box.pack_start(button,False,True,0)
        row.add(aux_box)

        window.connect("delete-event",_cancel_action)
        listbox.show_all()
        window.show_all()


    def make_usage_bar(self,usage):
        """ creates usage bar out of unicode characters""" 
        fill = float(usage) 
        fill = int(fill/10)
        # \u25A7 and \u25A1; maybe 2588
        #return u'\u2593'*fill + u'\u2591'*(10-fill)
        return u'\u2588'*fill + u'\u2592'*(10-fill)

    def get_human_readable(self,size):
         """ converts size in bytes to 
         human readable value in powers of 1024.
         Essentially, same as what df -h gives, or
         """
         prefix = ['B','KiB','MiB','GiB',
                   'TiB','PiB','EiB',
                   'ZiB','YiB'
         ]
         counter = 0
         while size/1024 > 0.9:
               counter = counter + 1
               size = size/1024
    
         return str(round(size,2)) +"  "+prefix[counter]



    def pad_string(self,orig,length,pad):
        """ Pads characters to string to make it 
            have specific length and returns it"""
        return orig + pad*(length - len(orig))

    def send_notif(self,note_obj,title,text):
        note_obj.update(title,text)
        note_obj.show()

    def mount_partition(self,*args):
        # TODO: implement error checking for mounting
        try:
            cmd = ['udisksctl','mount','-b','/dev/'+args[-1]]
            subprocess.check_output(cmd,stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            self.send_notif(self.note,'Mounting Error',e.output.decode())
            print(str(e))
            pass

    def get_filesystem_type(self,dev):
        args = ['system',
                'org.freedesktop.UDisks2',
                '/org/freedesktop/UDisks2/block_devices/' + dev,
                'org.freedesktop.UDisks2.Block',
                'IdType'
        ]
        try:
            return str(self.get_dbus_property(*args))
        except Exception as e:
            print(str(e))
            return None

    def get_bus_type(self,drive):
        """ Determines bus type of the bus via
        which a drive is connected """
        return self.get_dbus_property(
                   'system','org.freedesktop.UDisks2',
                   '/org/freedesktop/UDisks2/drives/' + drive,
                   'org.freedesktop.UDisks2.Drive','ConnectionBus'
        )
            
    def get_mountpoint_usage(self,mountpoint):
        """ performs statvfs syscall and calculates usage
        from the amount of filesystem blocks used. Returns
        tuple of percentage and actual byte size used"""
        fs = statvfs(mountpoint)
        used_blocks = float(fs.f_blocks) - float(fs.f_bfree)
        used_bytes = int(fs.f_bsize * used_blocks)
        pcent_use = 100 * (used_blocks/float(fs.f_blocks))
        return (str("{0:.2f}".format(pcent_use)),
                str(used_bytes) 
        )

    def get_partitions(self):
        contents = [ 
            'system','org.freedesktop.UDisks2','/org/freedesktop/UDisks2', 
            'org.freedesktop.DBus.ObjectManager','GetManagedObjects',None
        ]
        objects = self.get_dbus(*contents)
        
        partitions = []
        for item in objects:
            try:
                if 'block_devices'  in str(item):
                       contents = [
                           'system','org.freedesktop.UDisks2',item,
		           'org.freedesktop.UDisks2.Block','Drive'
                       ]
                       drive = self.get_dbus_property(*contents)
                       if drive == '/': continue
                      
                       mountpoint = self.get_mountpoint(item)
                       if not mountpoint: continue
                       mountpoint = mountpoint.replace('\x00','')
    
                       drive = str(drive).split('/')[-1]
                       pcent_use,byte_use = self.get_mountpoint_usage(mountpoint)

                       part = str(item.split('/')[-1])
                       partitions.append((part,drive,mountpoint,pcent_use,byte_use))                       

            except:
                pass
        partitions.sort()
        return partitions

    def get_mountpoint(self,dev_path):
        try:
            contents = [
                'system','org.freedesktop.UDisks2',dev_path,
                'org.freedesktop.UDisks2.Filesystem','MountPoints'
            ]
            data = self.get_dbus_property(*contents)[0]
    
        except:
            return None
        else:
            if len(data) > 0:
                return ''.join([ chr(byte) for byte in data])


    def get_unmounted_partitions(self):
        contents = [
            'system','org.freedesktop.UDisks2', 
            '/org/freedesktop/UDisks2','org.freedesktop.DBus.ObjectManager',
            'GetManagedObjects',None
        ]
        objects = self.get_dbus(*contents)
        
        partitions = []
        for item in objects:
            try:
                if 'block_devices'  in str(item):
                       drive = self.get_dbus_property('system',
                                        'org.freedesktop.UDisks2',
                                        item,
                                        'org.freedesktop.UDisks2.Block',
                                        'Drive')
                       if drive == '/': continue
                      
                       mountpoint = self.get_mountpoint(item)
                       if  mountpoint: continue

                       drive = str(drive).split('/')[-1]
                       part = str(item.split('/')[-1])
                       if not part[-1].isdigit(): continue
                       partitions.append((part,drive))                       
                       #print(partitions)

            except Exception as e:
                #print(e)
                pass

        partitions.sort()
        return partitions

    def get_dbus(self,bus_type,obj,path,interface,method,arg):
        """ utility: executes dbus method on specific interface"""
        if bus_type == "session":
            bus = dbus.SessionBus() 
        if bus_type == "system":
            bus = dbus.SystemBus()
        proxy = bus.get_object(obj,path)
        method = proxy.get_dbus_method(method,interface)
        if arg:
            return method(arg)
        else:
            return method()
    
    def get_dbus_property(self,bus_type,obj,path,iface,prop):
        """ utility:reads properties defined on specific dbus interface"""
        if bus_type == "session":
           bus = dbus.SessionBus()
        if bus_type == "system":
           bus = dbus.SystemBus()
        proxy = bus.get_object(obj,path)
        aux = 'org.freedesktop.DBus.Properties'
        props_iface = dbus.Interface(proxy,aux)
        props = props_iface.Get(iface,prop)
        return props

    def make_alias(self,*args):
        """ writes to config file user-defined alias"""
        partitions = [ i[0] for i in self.get_partitions() ]

        combo_values = '|'.join(partitions)
        command=[ 'zenity','--forms','--title',
                  'Make Alias','--text','',
                  '--add-combo','Partition','--combo-values',
                  combo_values,'--add-entry','Alias'    ]        
        user_input = self.run_cmd(command)
        if not user_input: return

        alias = user_input.decode().strip().split('|')

        existing_values = None
        if os.path.isfile(self.config_file):
            with open(self.config_file) as conf_file:
                try:
                    existing_values = json.load(conf_file)
                except ValueError:
                    pass
                    

        with open(self.config_file,'w') as conf_file:
             if existing_values:
                 existing_values[alias[0]] = alias[1]
             else:
                 existing_values = {alias[0]:alias[1]}
             json.dump(existing_values,conf_file,indent=4,sort_keys=True)
        

    def find_alias(self,part):
        """ searches the config file for partition alias"""
        if os.path.isfile(self.config_file):
            with open(self.config_file) as conf_file:
                try:
                    aliases = json.load(conf_file)
                except ValueError:
                    pass
                else:
                    if part in aliases:
                       return aliases[part]
                    else:
                       return None

    def process_autostart_file(self,*args):
        desktop_file = '.config/autostart/udisks-indicator.desktop'
        full_path = os.path.join(
                    self.user_home,
                    desktop_file
        )

        if self.prefs['autostart'] and not os.path.exists(full_path):
           original = '/usr/share/applications/udisks-indicator.desktop'
           if os.path.exists(original):
               shutil.copyfile(original,full_path)
        elif not self.prefs['autostart'] and os.path.exists(full_path):
           os.unlink(full_path)


    def open_mountpoint(self,*args):
        pid = subprocess.Popen(['xdg-open',args[-1]]).pid

    def open_disks_utility(self,*args):
        pid = subprocess.Popen(['gnome-disks']).pid

    def gsettings_get(self, schema, path, key):
        """utility: get value of gsettings schema"""
        if path is None:
            gsettings = Gio.Settings.new(schema)
        else:
            gsettings = Gio.Settings.new_with_path(schema, path)
        return gsettings.get_value(key)

    def run_cmd(self, cmdlist):
        """ utility: reusable function for running external commands """
        new_env = dict(os.environ)
        new_env['LC_ALL'] = 'C'
        try:
            stdout = subprocess.check_output(cmdlist, env=new_env)
        except subprocess.CalledProcessError:
            pass
        else:
            if stdout:
                return stdout

    def run(self):
        """ Launches the indicator """
        try:
            gtk.main()
        except KeyboardInterrupt:
            pass

    def quit(self, *args):
        """ closes indicator """
        gtk.main_quit()

def main():
    """ defines program entry point """
    indicator = UdisksIndicator()
    indicator.run()

if __name__ == '__main__':
    main()
