#!/usr/bin/python3
#
# (C) 2022 Guido Berhoerster <guido+ubports.com@berhoerster.name>
#
# This program 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.
#
# This package 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 this program. If not, see <http://www.gnu.org/licenses/>.
#
# On Debian systems, the complete text of the GNU General
# Public License version 3 can be found in "/usr/share/common-licenses/GPL-3".

# Acquire a display-on lock when repowerd has started, and release it when
# lightdm has started. Note that after releasing the display-on lock we get an
# additional 30s (at least) before the display turns off, leaving enough time
# for lightdm to fully start the session.
#
# The display-on lock is required so that the phone doesn't suspend during boot
# if long-running jobs are executed before the session starts (e.g. apparmor
# profile recompilation during the first boot after an upgrade, LP: #1623853).

import argparse
import logging
import signal
import sys

import dbus
from dbus.mainloop.glib import DBusGMainLoop
from gi.repository import GLib


DBUS_PROPERTIES_INTERFACE = "org.freedesktop.DBus.Properties"
DBUS_SERVICE_UNKNOWN_ERROR = "org.freedesktop.DBus.Error.ServiceUnknown"

USCREEN_NAME ="com.canonical.Unity.Screen"
USCREEN_PATH = "/com/canonical/Unity/Screen"
USCREEN_INTERFACE = "com.canonical.Unity.Screen"

SYSTEMD_NAME = "org.freedesktop.systemd1"
SYSTEMD_PATH = "/org/freedesktop/systemd1"
SYSTEMD_MANAGER_INTERFACE = "org.freedesktop.systemd1.Manager"
SYSTEMD_UNIT_INTERFACE = "org.freedesktop.systemd1.Unit"


class DisplayPowersaveBlocker:
    DEFAULT_TIMEOUT = 5 * 60 # s

    def __init__(self, timeout=0):
        self.bus = dbus.SystemBus()
        self.systemd1_proxy = None
        self.uscreen_proxy = None
        self.cookie = None
        self.timeout = timeout if timeout else self.DEFAULT_TIMEOUT
        self.loop = GLib.MainLoop()

    def keep_display_on(self):
        if self.cookie is not None:
            return

        try:
            self.cookie = (
                self.uscreen_proxy.keepDisplayOn(
                    dbus_interface=USCREEN_INTERFACE
                )
            )
        except dbus.exceptions.DBusException as e:
            logging.warning(f"failed to disable display powersaving")
        else:
            logging.debug(f"disabled display powersaving, cookie value: {self.cookie}")

    def allow_display_off(self):
        if self.cookie is None:
            return

        try:
            self.uscreen_proxy.removeDisplayOnRequest(self.cookie,
                    dbus_interface=USCREEN_INTERFACE)
        except dbus.exceptions.DBusException as e:
            if e.get_dbus_name() == DBUS_SERVICE_UNKNOWN_ERROR:
                logging.warning(f"repowerd is not running")
            else:
                raise e
        else:
            logging.debug("re-enabled display powersaving")

    def on_properties_changed(self, iface_name, changed_props,
            invalidated_props):
        logging.debug(
            f"unit properties changed: {', '.join(changed_props.keys())}"
        )
        if changed_props.get("ActiveState", "") in ("activating", "active"):
            logging.debug(f"unit is activating or active, quitting")
            self.loop.quit()

    def start_watching(self):
        unit = self.systemd1_proxy.LoadUnit("lightdm.service",
                dbus_interface=SYSTEMD_MANAGER_INTERFACE)
        logging.debug(f"loaded unit {unit}")
        unit_proxy = self.bus.get_object(SYSTEMD_NAME, unit)

        self.bus.add_signal_receiver(self.on_properties_changed,
                signal_name="PropertiesChanged", bus_name=SYSTEMD_NAME,
                path=unit)
        logging.debug(f"monitoring unit {unit} property changes")

        active = unit_proxy.Get(SYSTEMD_UNIT_INTERFACE, "ActiveState",
                dbus_interface=DBUS_PROPERTIES_INTERFACE)
        logging.debug(f"unit {unit} active state is {active}")
        if active in ("activating", "active"):
            logging.debug(f"unit is activating or active, quitting")
            self.loop.quit()

        self.keep_display_on()
        if self.cookie is None:
            self.loop.quit()

        return False

    def on_timeout(self):
        logging.debug(f"timed out, quitting")
        self.loop.quit()
        return False

    def on_signal(self):
        logging.debug(f"signal received, quitting")
        self.loop.quit()

    def run(self):
        try:
            self.uscreen_proxy = self.bus.get_object(USCREEN_NAME, USCREEN_PATH)
        except dbus.exceptions.DBusException as e:
            if e.get_dbus_name() == DBUS_SERVICE_UNKNOWN_ERROR:
                logging.warning(f"repowerd is not running, quitting")
                return
            else:
                raise e
        self.systemd1_proxy = self.bus.get_object(SYSTEMD_NAME, SYSTEMD_PATH)
        self.systemd1_proxy.Subscribe(dbus_interface=SYSTEMD_MANAGER_INTERFACE)

        GLib.idle_add(self.start_watching)
        GLib.timeout_add_seconds(self.timeout, self.on_timeout)
        GLib.unix_signal_add(GLib.PRIORITY_HIGH, signal.SIGINT, self.on_signal)
        GLib.unix_signal_add(GLib.PRIORITY_HIGH, signal.SIGTERM, self.on_signal)

        logging.debug(f"starting main loop")
        try:
            self.loop.run()
        finally:
            self.allow_display_off()


def timeout(value):
    value = int(value)
    if 0 < value < 60 * 60:
        return value
    raise ArgumentTypeError(f"Invalid timeout: {value}")


def main():
    aparser = argparse.ArgumentParser(
        description="block display powersaving until display manager is running"
    )
    aparser.add_argument("--timeout", action="store", type=timeout, default=0)
    aparser.add_argument("--debug", action="store_true", default=False)
    args = aparser.parse_args()

    DBusGMainLoop(set_as_default=True)

    logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s",
            level=logging.DEBUG if args.debug else logging.INFO)

    dpb = DisplayPowersaveBlocker(timeout=args.timeout)
    dpb.run()


if __name__ == "__main__":
    main()
