#!/usr/bin/python3
import sys
if len(sys.argv) > 1:
  print( sys.argv[0], '- tray applet to monitoring base parameters of Barium OS\n'
'and managing saves, please start it with no args, or by desktop file')
  quit()

import gi, os, subprocess, time, pickle, locale
gi.require_version('Gtk', '3.0')
gi.require_version('AppIndicator3', '0.1')
from gi.repository import Gtk, GLib
from gi.repository import AppIndicator3 as appindicator

from locale import gettext as _

with open('/proc/cmdline', 'r') as f:
  cmdline = f.readline()

with open('/proc/mounts', 'r') as f:
  mounts = f.readlines()

with open('/etc/initvars', 'r') as f:
  initvars = {}
  for line in f.read().splitlines():
    initvars[line.split('=')[0]] = line.split('=', 1 )[1]

isbackup = False
if 'uird.from=/ROSA-SYSTEM.bak' in cmdline:
  isbackup = True

ischanges = False
if initvars['uird_mode'] == 'changes' in cmdline: 
  ischanges = True 

isluks = False
for f in os.listdir('/dev/mapper'):
  for line in mounts:
    if '/dev/mapper/' + f  in line:
      isluks = True
      break
  if isluks:
    break

istoxzm = False
if initvars['uird_mode'] == 'toxzm' and os.path.exists('/run/initramfs/remount'):
  istoxzm = True

ishomexzm = False
if os.path.exists(initvars['SYSMNT'] + '/layer-base/1/saves/homes.xzm') and istoxzm:
  ishomexzm = True

APP = 'barium_applet'
LOCALE_DIR = '/usr/share/locale'

##### test block ##########
#LOCALE_DIR = '../../locale'
###########################

locale.bindtextdomain(APP, LOCALE_DIR)
locale.textdomain(APP)
h = {}
h['ROOTFS'] = _('ROOTFS')
h['DISK_USAGE'] = _('DISK USAGE')
h['MEM'] = _('MEM')
h['LAST_SAVE'] = _('LAST SAVE')
h['HOME_DIR_SIZE'] = _('HOME DIR SIZE')
h['HOME_MODULE_SIZE'] = _('HOME MODULE SIZE')
h['SAVE_TIMEOUT'] = _('SAVE TIMEOUT')
used_ = _('used')
available_ = _('available')
min_ = _('min')
Mb = _('Mb')
disabled_ = _('disabled')
maxlen = len(max( h.values() , key=len)) + 2
for key, val in h.items():
  quant = (maxlen - len(val) + 2) // 5
  tabs = str('_' * quant).replace('_', '\t')
  h[key] = val + ': ' + tabs  

#default values
config = {}
config['times'] = (1, 5, 10, 30, disabled_ ) # min
config['synctimeout'] = 30 # min (one of times)
config['upd_times'] = (1, 3, 5, 10, 30) # sec
config['timeout'] = 5 # sec (one of upd_times)
 
new_interval = False

def pickle_write():
  with open(cfgdir + '/dump', 'wb') as f:
   pickle.dump(config, f)

cfgdir = '/home/' + os.environ['USER'] + '/.config/barium_applet'
if os.path.exists(cfgdir + '/dump'):
  with open(cfgdir + '/dump', 'rb') as f:
    config = pickle.load(f)
else:
  os.makedirs(cfgdir, exist_ok=True)
  pickle_write()

def fn_exit(w, data):
  Gtk.main_quit()

def fn_check_updates(w, data):
  subprocess.call( [ 'barium', 'update' ]) 

def fn_tobackup(w, data):
  script = '''#!/bin/bash
DISTRDEV=$(grep -m1 " %s/layer-base/0 " /proc/self/mountinfo | awk '{ print $10 }') 
DISTRMPOINT=$(findmnt -lnf $DISTRDEV -o TARGET)
DISTRPATH=$DISTRMPOINT$(grep -m1 " %s/layer-base/0 " /proc/self/mountinfo | awk '{ print $4 }')
if [ "${DISTRPATH##*.}" != bak ] ; then
	zenity --error --text "%s"
	exit
fi 
mount -o remount,rw $DISTRMPOINT
mv ${DISTRPATH/.bak} ${DISTRPATH/.bak}.upd
mv ${DISTRPATH} ${DISTRPATH/.bak}
if [ $? == 0 ] ; then
	zenity --info --text "%s"
else
	zenity --info --text "%s"
fi
''' % ( initvars['SYSMNT'], initvars['SYSMNT'], _('Not a backup'), _('System was switched to boot from backup. Please reboot'), _('Unknown error, sorry'))
  with open('/tmp/switch_to_backup.sh', 'w') as f:
    print( script, file=f )
  os.chmod('/tmp/switch_to_backup.sh', 0o755) 
  subprocess.call( [ 'bashsu', '/tmp/switch_to_backup.sh' ] )

def fn_toggle_saves(w, data):
  if os.path.exists('/run/initramfs/shutdown'):
    ret = subprocess.call( [ 'bashsu', 'mv', '/run/initramfs/shutdown' , '/run/initramfs/shutdown.bak' ])
    if not ret:
      menu_toggle_saves.set_label(_('Enable saves'))
  else:
    ret = subprocess.call( [ 'bashsu', 'mv', '/run/initramfs/shutdown.bak' , '/run/initramfs/shutdown' ])
    if not ret: 
      menu_toggle_saves.set_label(_('Disable saves'))

# Генератор функций для barium marriage
fn_marriage = {}
def f3(key):
  subprocess.call( [ 'bashsu', 'barium', 'marriage', key ] )
for i in ( 'rmCurrent', 'addKey', 'rmAll' ) :
  fn_marriage[i] = lambda x, y, i=i: f3(i) 

def fn_rmbak(w, data):
  subprocess.call( [ 'bashsu', 'rm', '-f', initvars['SYSMNT'] + '/layer-base/1/saves/*.xzm.bak' ]) 

def fn_synchome(*args):
  ret = subprocess.call(['synchome.x'],)
  if not ret:
    with open( '/tmp/synchome', 'w' ) as f:
      f.write(str(int(time.time())))

#генератор функций для каждого значение в кортеже times
fn_save = {}
def f(time_):
  global config
  if type(time_) == int:
    fn_synchome()
  if type(config['synctimeout']) == int:
    submenu_save_item[config['synctimeout']].set_label( str(config['synctimeout']) + ' ' + _('min') )
  else:
    submenu_save_item[config['synctimeout']].set_label( str(config['synctimeout']) )
  config['synctimeout'] = time_  
  pickle_write()
for i in config['times']:
  fn_save[i] = lambda x, y, i=i: f(i) 

#генератор функций для каждого значение в кортеже upd_times
fn_icon_update = {}
def f2(time_):
  global config
  global timer_id
  global new_interval
  submenu_icon_item[config['timeout']].set_label(str(config['timeout']) + ' ' + _("sec") )
  config['timeout'] = time_
  new_interval = True
  submenu_icon_item[config['timeout']].set_label(str(config['timeout']) + ' ' + _("sec") + _(" *") ) 
  new_interval = True 
  pickle_write()
for i in config['upd_times']:
  fn_icon_update[i] = lambda x, y, i=i: f2(i) 


def fn_getvalues(ind_app):
  global new_interval
  available = {}
  used = {}
  for name, path in ([ 'rootfs',  '/' ], [ 'disk', initvars['SYSMNT'] + '/layer-base/1' ]):
    with subprocess.Popen([ 'df', path, '-m', '/'], stdout=subprocess.PIPE) as fs:
      lines = fs.stdout.readlines()
      valueline = lines[1].decode('utf-8').split()
      available[name] = valueline[3]
      used[name] = valueline[4][0:-1]
  with subprocess.Popen(['free', '-m'], stdout=subprocess.PIPE) as free:
    lines = free.stdout.readlines()
    valueline = lines[1].decode('utf-8').split()
    available['free'] = valueline[6]
    used['free'] = float('{:.2f}'.format(int(valueline[2]) * 100 / int(valueline[1])))
  if os.path.exists('/tmp/synchome'):
    with open( '/tmp/synchome', 'r' ) as f:
      lastsave = str(f.readlines()[0])
      timestamp = int(time.time())
      delta = int((timestamp - int(lastsave)) / 60)
      if type(config['synctimeout']) == int  and delta >= config['synctimeout']:
        fn_synchome()
  else:
      delta = _('never')
  if ishomexzm:
    homesize = float('{:.2f}'.format(os.stat(initvars['SYSMNT'] + '/layer-base/1/saves/homes.xzm').st_size / 1024 / 1024))
  else:
    homesize = 0
  with subprocess.Popen([ 'du', '/home', '-sm'], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) as home:
    homedirsize = str(home.stdout.readline().decode('utf-8').split()[0])
  if type(config['synctimeout']) == int:
    timeoutstr = str(config['synctimeout']) + ' ' + min_
  else:
    timeoutstr = config['synctimeout']
  title = f'''
{h['ROOTFS']}{available_}: {available['rootfs']}{Mb}, {used_}: {used['rootfs']}%
{h['DISK_USAGE']}{available_}: {available['disk']}{Mb}, {used_}: {used['disk']}%
{h['MEM']}{available_}: {available['free']}{Mb}, {used_}: {used['free']}%
{h['HOME_DIR_SIZE']}{homedirsize}{Mb}
'''
  if ishomexzm:
    title += f'''
{h['HOME_MODULE_SIZE']}{homesize}{Mb}'
{h['LAST_SAVE']}{delta} {min_}
{h['SAVE_TIMEOUT']}{timeoutstr}
'''
  ind_app.set_title(title)
  
  color = 'green'
  if ishomexzm:
    if int(homedirsize) / 1.5 >= int(available['disk']) or \
    int(homesize) >= 500:
      color = 'yellow'
  
  if not ischanges and int(used['rootfs']) >= 80:
    color = 'yellow'
  
  if int(used['free']) >= 80:
    color = 'yellow'

  if ishomexzm:
    if int(homedirsize) / 2 >= int(available['disk']) or \
    int(homesize) >= 1000:
      color = 'red'
      
  if istoxzm and not os.path.exists('/run/initramfs/shutdown'):
    color = 'red'
    
  if not ischanges and int(used['rootfs']) >= 95:
    color = 'red'
  
  if int(used['free']) >= 95:
    color = 'red'
  


  if color == 'red':
    if ind_app.get_icon != 'error':
      ind_app.set_icon_full('error', 'error')
      menu_rmback.show()
  elif color == 'yellow':
    if ind_app.get_icon != 'warning':
      ind_app.set_icon('warning', 'warning')
      menu_rmback.show()
  elif color == 'green':
    if ind_app.get_icon != 'ok':
      ind_app.set_icon_full('ok', '')
      menu_rmback.hide()
  if ishomexzm:
    if type(config['synctimeout']) == int:
      submenu_save_item[config['synctimeout']].set_label( str(config["synctimeout"]) + ' ' + _('min') + ' *' )
    else:
      submenu_save_item[config['synctimeout']].set_label(disabled_ + ' *')

  if new_interval:
    new_interval = False
    GLib.timeout_add(config['timeout'] * 1000, fn_getvalues, ind_app)
    return False
  else:
    return True

submenu_save_item = {}
def show_saves(times):
  for itime  in times:
    if type(itime) == int:
      submenu_save_item[itime] = Gtk.MenuItem( label=str(itime) + ' ' + _('min') )
    else:
      submenu_save_item[itime] = Gtk.MenuItem( label=str(itime))
    submenu_save_item[itime].connect("activate", fn_save[itime], '')
    submenu_save_item[itime].show()
    submenu_save.append(submenu_save_item[itime])

submenu_icon_item = {}
def show_icon_update(times):
  for itime  in times:
    submenu_icon_item[itime] = Gtk.MenuItem( label=str(itime) + ' ' + _('sec') )
    submenu_icon_item[itime].connect("activate", fn_icon_update[itime], '')
    submenu_icon_item[itime].show()
    submenu_icon_update.append(submenu_icon_item[itime])

def show_marriage(key, label):
  submenu_marriage_item[key] = Gtk.MenuItem( label=label )
  submenu_marriage_item[key].connect("activate", fn_marriage[key], '')
  submenu_marriage_item[key].show()
  submenu_marriage_menu.append(submenu_marriage_item[key])

ind_app = appindicator.Indicator.new_with_path (
   "barium-status-indicator",
   "ok",
   appindicator.IndicatorCategory.APPLICATION_STATUS,
   '/usr/share/rosa-rw/barium_applet')
ind_app.set_status (appindicator.IndicatorStatus.ACTIVE)
# create a menu
menu = Gtk.Menu()

if ishomexzm:
  menu_save_now = Gtk.MenuItem(label=_('Save home now'))
  menu_save_now.connect("activate", fn_synchome, '')
  menu_save_now.show()
  menu.append(menu_save_now)

  ######## autosave
  submenuItem = Gtk.MenuItem(label=_('Autosave'))
  submenu_save = Gtk.Menu()
  submenuItem.set_submenu(submenu_save)
  menu.append(submenuItem)
  submenuItem.show()
  show_saves(config['times'])

######## update applet
submenuItem2 = Gtk.MenuItem(label=_('Applet update interval'))
submenu_icon_update = Gtk.Menu()
submenuItem2.set_submenu(submenu_icon_update)
menu.append(submenuItem2)
submenuItem2.show()
show_icon_update(config['upd_times'])
submenu_icon_item[config['timeout']].set_label(str(config['timeout']) + ' ' + _('sec') + ' *' )

######## updates
menu_updates = Gtk.MenuItem(label=_('Check updates'))
menu_updates.connect("activate", fn_check_updates, '')
menu_updates.show()
menu.append(menu_updates)

if isluks:
  ######## marriage menu
  submenuItem3 = Gtk.MenuItem(label=_('Known machines'))
  submenu_marriage_menu = Gtk.Menu()
  submenuItem3.set_submenu(submenu_marriage_menu)
  menu.append(submenuItem3)
  submenuItem3.show()
  submenu_marriage_item = {}
  show_marriage('addKey', _('add machine to list of known') )
  show_marriage('rmCurrent', _('remove this machine from list of known') )
  show_marriage('rmAll', _('remove all machines from list') )

if istoxzm:
  ######## saves toggle
  if os.path.exists('/run/initramfs/shutdown'):
    toggle_label = _('Disable saves')
  elif os.path.exists('/run/initramfs/shutdown.bak'):
    toggle_label = _('Enable saves')
  else:
    toggle_label = None

  if toggle_label:
    menu_toggle_saves = Gtk.MenuItem(label=toggle_label)
    menu_toggle_saves.connect("activate", fn_toggle_saves, '')
    menu_toggle_saves.show()
    menu.append(menu_toggle_saves)

######### delete backup mods
menu_rmback = Gtk.MenuItem(label=_('Delete backup modules'))
menu_rmback.connect("activate", fn_rmbak, '')
menu.append(menu_rmback)

if isbackup:
  ######### switch to backup
  menu_tobackup = Gtk.MenuItem(label=_('Switch to backup'))
  menu_tobackup.connect("activate", fn_tobackup, '')
  menu_tobackup.show()
  menu.append(menu_tobackup)

######### separator
separ = Gtk.SeparatorMenuItem()
menu.append(separ)
separ.show()

######### exit
menu_exit = Gtk.MenuItem(label=_('Exit'))
menu_exit.connect("activate", fn_exit, '')
menu_exit.show()
menu.append(menu_exit)

#########
menu.show()
ind_app.set_menu(menu)
timer_id = GLib.timeout_add(config['timeout'] * 1000, fn_getvalues, ind_app)
Gtk.main()
