#!/usr/bin/env python3

import sys
import dbus
import re
import datetime
import subprocess
# import pprint


class AutoSuspender(object):

    def __init__(self):
        self.__bus = dbus.SystemBus()
        self.__timer_re = re.compile('^syncoid-pull-.*\.timer$')
        self.__zpool_scrub_in_progress_re = re.compile('^\s*scan:.*scrub in progress.*$')
        self.__zpool_resilver_in_progress_re = re.compile('^\s*scan:.*resilver in progress.*$')
        # self.__pp = pprint.PrettyPrinter(indent=2)

    def _get_interface(self, dest, object, interface):
        try:
            obj = self.__bus.get_object(dest, object)
            return dbus.Interface(obj, interface)
        except dbus.exceptions.DBusException as error:
            print("ERR: %s" % error)
            sys.exit(1)

    def _get_logind_manager_interface(self):
        return self._get_interface("org.freedesktop.login1", "/org/freedesktop/login1", "org.freedesktop.login1.Manager")

    def _get_systemd_manager_interface(self):
        return self._get_interface("org.freedesktop.systemd1", "/org/freedesktop/systemd1", "org.freedesktop.systemd1.Manager")

    def _get_systemd_timer_properties(self, path):
        return self._get_interface("org.freedesktop.systemd1", path, "org.freedesktop.DBus.Properties").GetAll("org.freedesktop.systemd1.Timer")

    def _get_users_logged_in(self):
        manager = self._get_logind_manager_interface()
        users = []
        for user in manager.ListUsers():
            uid = int(user[0])
            name = str(user[1])
            users.append((name, uid))
        return users

    def _get_syncoid_timers(self):
        manager = self._get_systemd_manager_interface()
        timers = []
        for unit in manager.ListUnits():
            name = str(unit[0])
            path = str(unit[6])
            if not self.__timer_re.match(name):
                continue
            timers.append((name, path))
        return timers

    def _get_timer_next_elapse(self, path):
        props = self._get_systemd_timer_properties(path)
        try:
            return datetime.datetime.fromtimestamp(int(props['NextElapseUSecRealtime']) / 1000000)
        except ValueError:
            return None

    def _get_zpools(self):
        pools = subprocess.run(["zpool", "list", "-H", "-o", "name"], check=True, stdout=subprocess.PIPE).stdout.splitlines()
        return [pool.decode('utf-8') for pool in pools]

    def _is_zpool_scrub_or_resilver_in_progress(self, pool):
        lines = subprocess.run(["zpool", "status", pool], check=True, stdout=subprocess.PIPE).stdout.splitlines()
        for line in lines:
            if self.__zpool_scrub_in_progress_re.match(line.decode('utf-8')):
                return True
            if self.__zpool_resilver_in_progress_re.match(line.decode('utf-8')):
                return True
        return False

    def check(self):
        result = True

        users = self._get_users_logged_in()
        if(len(users) > 0):
            print(" [NO]  %d users logged in: %s" % (len(users), ", ".join([user[0] for user in users])))
            result = False
        else:
            print(" [ok]  no users logged in.")

        timers = self._get_syncoid_timers()
        for timer in timers:
            next = self._get_timer_next_elapse(timer[1])
            if next == None:
                print(" [NO]  the next elapse time of '%s' is invalid, assuming the unit is currently active" % (timer[0]))
                result = False
                continue

            until = next - datetime.datetime.now()
            if(until < datetime.timedelta(minutes=10)):
                print(" [NO]  timer '%s' elapses in less than 10 minutes -> %s (%s)" % (timer[0], next, until))
                result = False
            else:
                print(" [ok]  timer '%s' elapses in %s (%s)" % (timer[0], until, next))

        zpools = self._get_zpools()
        for pool in zpools:
            if self._is_zpool_scrub_or_resilver_in_progress(pool):
                print(" [NO]  zpool scrub/resilver is in progress on pool '%s'" % (pool))
                result = False
            else:
                print(" [ok]  zpool '%s' is not scrubbing/resilvering" % (pool))

        return result

    def suspend_system(self):
        print("trying to suspend system")
        manager = self._get_logind_manager_interface()
        if not manager.CanSuspend():
            print("ERR: system can not be suspended")
            return
        try:
            manager.Suspend(False)
        except dbus.exceptions.DBusException as error:
            print("ERR: %s" % error)

    def _reformat_hdparm_output(self, lines):
        result = []
        entry = ""
        for line in lines:
            line = line.decode('utf-8')
            if line == '':
                if entry != '':
                    result.append(entry)
                entry = ""
                continue
            if line.strip().startswith("drive state is:"):
                line = ' ' + line.split(':', 2)[1].strip()
            entry += line
        result.append(entry)
        return result

    def spindown_disks(self, disks):
        print("current disk power states:")
        lines = subprocess.run(["hdparm", "-C"] + disks, check=True, stdout=subprocess.PIPE).stdout.splitlines()
        for line in self._reformat_hdparm_output(lines):
            print("  %s" % line)
        print("issuing standby commands:")
        lines = subprocess.run(["hdparm", "-y"] + disks, check=True, stdout=subprocess.PIPE).stdout.splitlines()
        for line in self._reformat_hdparm_output(lines):
            print("  %s" % line)


def print_usage():
    print("usage: %s (system|disks) [ /dev/sda [ /dev/sdb [ ... ]]] " % sys.argv[0])


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print_usage()
        sys.exit(1)
    mode = sys.argv[1]
    disks = []
    if mode not in ["system", "disks"]:
        print_usage()
        sys.exit(1)
    if mode == "disks":
        if len(sys.argv) < 3:
            print_usage()
            sys.exit(1)
        else:
            disks = sys.argv[2:]

    s = AutoSuspender()
    print("checking for suspend inhibitors:")
    if s.check():
        if mode == "system":
            s.suspend_system()
        elif mode == "disks":
            s.spindown_disks(disks)
    else:
        print("not suspending because at least one check failed.")