#!/usr/bin/python3
import importlib.machinery
import importlib.util
import os
from pathlib import Path
import re
import subprocess
import sys
import threading
import uuid

import gi

gi.require_version("Gio", "2.0")
from gi.repository import Gio, GLib

NAME = "org.freedesktop.login1"
ROOT = "/org/freedesktop/login1"
INTERFACE = NAME + ".Manager"
UPOWER = "org.freedesktop.UPower"
UPOWER_ROOT = "/org/freedesktop/UPower"
NATIVE_PREPARE_MS = 1500
AGENT = Path(__file__).resolve().with_name("droiddeck-agent")
ASLEEP_WAIT = 6
INHIBIT_WAIT = 12
RELEASE_RETRIES = 8
INTROSPECTION = """
<node>
  <interface name="org.freedesktop.login1.Manager">
    <method name="Suspend"><arg type="b" direction="in"/></method>
    <method name="SuspendWithFlags"><arg type="t" direction="in"/></method>
    <method name="CanSuspend"><arg type="s" direction="out"/></method>
    <property name="PreparingForSleep" type="b" access="read"/>
    <signal name="PrepareForSleep"><arg type="b"/></signal>
  </interface>
</node>
"""
UPOWER_INTROSPECTION = """
<node>
  <interface name="org.freedesktop.UPower">
    <method name="Suspend"/>
  </interface>
</node>
"""


class LoginManager:
    def __init__(self, bus, directory):
        self.bus = bus
        self.directory = Path(directory)
        self.request = self.directory / "steam-sleep"
        self.state = self.directory / "steam-sleep-state"
        self.wake = self.directory / "steam-wake"
        self.native_request = self.directory / "steam-native-sleep"
        self.native_ready = self.directory / "steam-native-ready"
        self.pending = None
        self.pending_idle = False
        self.waiting = set()
        self.timeout = None
        self.preparation = None
        self.inhibited = False
        self.inhibit_lock = threading.Lock()
        self.release = None
        self.generation = 0
        info = Gio.DBusNodeInfo.new_for_xml(INTROSPECTION).interfaces[0]
        self.registration = bus.register_object(ROOT, info, self.method, self.property, None)
        upower_info = Gio.DBusNodeInfo.new_for_xml(UPOWER_INTROSPECTION).interfaces[0]
        self.upower_registration = bus.register_object(UPOWER_ROOT, upower_info, self.method, None, None)
        self.monitor = Gio.File.new_for_path(str(self.directory)).monitor_directory(Gio.FileMonitorFlags.NONE, None)
        self.monitor.connect("changed", self.changed)

    def property(self, connection, sender, path, interface, name):
        if name == "PreparingForSleep":
            return GLib.Variant("b", self.pending is not None)
        return None

    def notify(self, sleeping):
        self.bus.emit_signal(None, ROOT, INTERFACE, "PrepareForSleep", GLib.Variant("(b)", (sleeping,)))
        self.bus.emit_signal(None, ROOT, "org.freedesktop.DBus.Properties", "PropertiesChanged",
                             GLib.Variant("(sa{sv}as)", (INTERFACE, {
                                 "PreparingForSleep": GLib.Variant("b", sleeping),
                             }, [])))
        print("== login1: " + ("pausing the session" if sleeping else "session resumed"), flush=True)

    def method(self, connection, sender, path, interface, name, args, reply):
        if name == "CanSuspend":
            reply.return_value(GLib.Variant("(s)", ("yes",)))
        elif name in ("Suspend", "SuspendWithFlags"):
            print("== login1: %s.%s asked by %s" % (interface.rsplit(".", 1)[-1], name, sender), flush=True)
            if self.pending:
                reply.return_dbus_error(NAME + ".OperationInProgress", "The session is already pausing")
                return
            self.pending = uuid.uuid4().hex
            self.pending_idle = False
            self.notify(True)
            reply.return_value(None)
            self.bus.flush(None, self.publish_after_flush, self.pending)

    def publish_after_flush(self, connection, result, token):
        try:
            connection.flush_finish(result)
        except GLib.Error as error:
            print("== login1: cannot complete suspend handshake: " + str(error), file=sys.stderr)
            self.resume()
            return
        if self.pending == token:
            self.hold(token, self.publish, self.request)

    def publish(self, request):
        self.check_state()
        if self.pending is None:
            return False
        try:
            staged = request.with_suffix(".tmp")
            staged.write_text(self.pending + "\n")
            staged.replace(request)
            self.timeout = GLib.timeout_add_seconds(5, self.check_ack)
        except OSError as error:
            print("== login1: cannot request pause: " + str(error), file=sys.stderr)
            self.resume()
        return False

    def prepare_native(self):
        try:
            words = self.native_request.read_text().split()
            self.native_request.unlink(missing_ok=True)
        except OSError:
            return
        token = words[0] if words else ""
        if not re.fullmatch(r"[a-f0-9]{32}", token) or words[1:] not in ([], ["idle"]) or self.pending is not None:
            return
        if self.read_state() == token + ":awake":
            return
        self.pending = token
        self.pending_idle = words[1:] == ["idle"]
        self.waiting = {"hold"} if self.pending_idle else {"hold", "prepare"}
        self.hold(token, self.prepared, token, "hold")
        if self.pending_idle:
            return
        self.notify(True)
        self.bus.flush(None, self.prepare_after_flush, token)

    def prepare_after_flush(self, connection, result, token):
        try:
            connection.flush_finish(result)
        except GLib.Error as error:
            print("== login1: cannot prepare native sleep: " + str(error), file=sys.stderr)
            self.resume()
            return
        self.check_state()
        if self.pending == token:
            self.preparation = GLib.timeout_add(NATIVE_PREPARE_MS, self.finish_preparation, token)

    def finish_preparation(self, token):
        self.preparation = None
        self.prepared(token, "prepare")
        return False

    def prepared(self, token, step):
        if self.pending != token:
            return False
        self.waiting.discard(step)
        if not self.waiting:
            self.publish(self.native_ready)
        return False

    def hold(self, token, then, *args):
        self.cancel_release()

        def work():
            self.set_inhibit(True)
            GLib.idle_add(self.after_hold, token, then, args)

        threading.Thread(target=work, daemon=True).start()

    def after_hold(self, token, then, args):
        if self.pending == token:
            then(*args)
        elif self.pending is None:
            self.release_later()
        return False

    def set_inhibit(self, inhibit, generation=None):
        with self.inhibit_lock:
            if (generation is not None and generation != self.generation) or self.inhibited == inhibit:
                return
            if steam_idle_inhibit(inhibit):
                self.inhibited = inhibit
                print("== login1: Steam idle sleep " + ("held" if inhibit else "released"), flush=True)
            else:
                self.inhibited = None

    def cancel_release(self):
        self.generation += 1
        if self.release is not None:
            GLib.source_remove(self.release)
            self.release = None

    def release_later(self):
        self.cancel_release()
        if self.inhibited is False:
            return
        generation = self.generation
        threading.Thread(target=lambda: GLib.idle_add(self.schedule_release, generation, steam_idle_suspend_seconds()),
                         daemon=True).start()

    def schedule_release(self, generation, seconds, attempt=0):
        if generation == self.generation and self.pending is None:
            self.release = GLib.timeout_add_seconds(seconds or 0, self.release_now, generation, attempt)
        return False

    def release_now(self, generation, attempt):
        self.release = None

        def work():
            self.set_inhibit(False, generation)
            if self.inhibited is not False and attempt < RELEASE_RETRIES:
                GLib.idle_add(self.schedule_release, generation, 1 << attempt, attempt + 1)

        threading.Thread(target=work, daemon=True).start()
        return False

    def changed(self, monitor, file, other, event):
        if file.get_basename() == self.native_request.name or (other and other.get_basename() == self.native_request.name):
            self.prepare_native()
        if file.get_basename() == self.wake.name or (other and other.get_basename() == self.wake.name):
            if self.wake.exists():
                self.wake.unlink(missing_ok=True)
                self.woken()
        if file.get_basename() == self.state.name or (other and other.get_basename() == self.state.name):
            self.check_state()

    def read_state(self):
        try:
            return self.state.read_text().strip()
        except OSError:
            return ""

    def check_state(self):
        if self.pending and self.read_state() == self.pending + ":awake":
            self.resume()

    def check_ack(self):
        self.timeout = None
        if self.pending and self.read_state() != self.pending + ":paused":
            self.request.unlink(missing_ok=True)
            self.native_ready.unlink(missing_ok=True)
            self.resume()
        return False

    def woken(self):
        if self.pending is not None:
            idle = self.pending_idle
            self.resume(force=True)
            if not idle:
                return
        threading.Thread(target=lambda: GLib.idle_add(self.replay_unless_awake, steam_asleep()), daemon=True).start()

    def replay_unless_awake(self, asleep):
        if asleep is False and self.pending is None:
            print("== login1: Steam is awake; no sleep/wake replay", flush=True)
            return False
        self.resume(force=True)
        return False

    def resume(self, force=False):
        if self.pending is None:
            if force:
                self.notify(True)
                self.bus.flush(None, self.finish_recovery, None)
            return
        idle = self.pending_idle
        self.pending = None
        self.pending_idle = False
        if self.preparation is not None:
            GLib.source_remove(self.preparation)
            self.preparation = None
        if self.timeout is not None:
            GLib.source_remove(self.timeout)
            self.timeout = None
        if not idle:
            self.notify(False)
        self.release_later()

    def finish_recovery(self, connection, result, _data):
        try:
            connection.flush_finish(result)
        except GLib.Error as error:
            print("== login1: cannot flush wake recovery: " + str(error), file=sys.stderr)
        if self.pending is None:
            self.notify(False)


def agent(*args, timeout=ASLEEP_WAIT):
    try:
        done = subprocess.run([str(AGENT), *args], stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                              stderr=subprocess.DEVNULL, text=True, timeout=timeout)
    except (OSError, subprocess.SubprocessError):
        return None, ""
    return done.returncode, done.stdout


def steam_asleep():
    return {0: True, 5: False}.get(agent("--steam-asleep")[0])


def steam_idle_inhibit(inhibit):
    return agent("--steam-idle-inhibit", "on" if inhibit else "off", timeout=INHIBIT_WAIT)[0] == 0


def steam_idle_suspend_seconds():
    code, output = agent("--steam-idle-suspend-seconds")
    return int(output) if code == 0 and output.strip().isdigit() else None


def load_netmanager():
    path = Path(__file__).resolve().with_name("droiddeck-netmanager")
    loader = importlib.machinery.SourceFileLoader("droiddeck_netmanager", str(path))
    spec = importlib.util.spec_from_loader(loader.name, loader)
    module = importlib.util.module_from_spec(spec)
    loader.exec_module(module)
    return module


def main():
    directory = os.environ.get("BL_LAUNCH_DIR")
    if not directory or not Path(directory).is_dir():
        sys.exit("droiddeck-login1: no DroidDeck session directory")
    ready = Path(directory) / "steam-sleep-ready"
    ready.unlink(missing_ok=True)
    loop = GLib.MainLoop()
    manager = []
    acquired_names = set()

    def connected(bus, name):
        if not manager:
            manager.append(LoginManager(bus, directory))

    def acquired(bus, name):
        acquired_names.add(name)
        if acquired_names == {NAME, UPOWER}:
            ready.write_text("ready\n")

    def lost(bus, name):
        ready.unlink(missing_ok=True)
        loop.quit()

    owners = [Gio.bus_own_name(Gio.BusType.SYSTEM, name, Gio.BusNameOwnerFlags.NONE, connected, acquired, lost)
              for name in (NAME, UPOWER)]
    if "--with-netmanager" in sys.argv[1:]:
        owners.append(load_netmanager().own(loop))
    try:
        loop.run()
    finally:
        ready.unlink(missing_ok=True)
        for owner in owners:
            Gio.bus_unown_name(owner)


if __name__ == "__main__":
    main()
