#!/usr/bin/python3
#
""" Launches or restarts cinnamon
"""

import os
import sys

# Force X11 before any display connection is made
os.environ["XDG_SESSION_TYPE"] = "x11"
os.environ["GDK_BACKEND"] = "x11"
os.environ["DISPLAY"] = os.environ.get("DISPLAY", ":0")

import time
import gettext
from setproctitle import setproctitle
from pathlib import Path
import psutil
import subprocess
import threading
import gi
import shutil
gi.require_version('Gtk', '3.0')  # noqa
from gi.repository import Gtk, GLib, Gio, GLib

FALLBACK_COMMAND = "metacity"
FALLBACK_ARGS = ("--replace",)

gettext.install("cinnamon", "/usr/share/locale")

panel_process_name = None
panel_cmd = None
if shutil.which("mate-panel"):
    panel_process_name = "mate-panel"
    panel_cmd = "mate-panel --replace &"
elif shutil.which("gnome-panel"):
    panel_process_name = "gnome-panel"
    panel_cmd = "gnome-panel --replace &"
elif shutil.which("tint2"):
    panel_process_name = "tint2"
    panel_cmd = "killall tint2; tint2 &"

# Used as a decorator to run things in the background
def async_function(func):
    def wrapper(*args, **kwargs):
        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
        thread.daemon = True
        thread.start()
        return thread
    return wrapper

# Used as a decorator to run things in the main loop, from another thread
def idle_function(func):
    def wrapper(*args):
        GLib.idle_add(func, *args)
    return wrapper

class Launcher:

    def __init__(self):
        self.fb_pid = 0
        self.kill_lock = threading.Lock()

        self.polkit_agent_proc = None
        self.nm_applet_proc = None

        subprocess.Popen(
            ["systemctl", "--user", "start", "cinnamon-session.service"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        self.cinnamon_pid = os.fork()
        if self.cinnamon_pid == 0:
            os.execvp("cinnamon", ("cinnamon", "--x11", "--replace", ) + tuple(sys.argv[1:]))
        else:
            self.settings = Gio.Settings(schema="org.cinnamon.launcher")
            self.settings.connect("changed::memory-limit-enabled", self.on_settings_changed)
            self.killed = False
            if self.settings.get_boolean("memory-limit-enabled"):
                print("Cinnamon memory limit enabled: %d MB" % self.settings.get_int("memory-limit"))
                self.monitor_memory()
            self.wait_for_process()
            Gtk.main()

    def on_settings_changed(self, settings, key):
        print("Memory profiler status changed, restarting Cinnamon.")
        self.restart_cinnamon()

    @async_function
    def wait_for_process(self):
        exit_status = os.waitpid(self.cinnamon_pid, 0)[1]
        if not os.WIFEXITED(exit_status):
            with self.kill_lock:
                if self.killed:
                    # Restart Cinnamon
                    self.restart_cinnamon()
                else:
                    self.fb_pid = os.fork()
                    if self.fb_pid == 0:
                        # Start the fallback panel
                        if panel_cmd is not None:
                            os.system(panel_cmd)
                        # Start the fallback window manager
                        os.execvp(FALLBACK_COMMAND, (FALLBACK_COMMAND,) + FALLBACK_ARGS)
                    else:
                        self.launch_fallback_helpers()
                        self.confirm_restart()
        else:
            Gtk.main_quit()

    def launch_fallback_helpers(self):
        agents = ("/usr/lib/policykit-1-gnome/polkit-gnome-authentication-agent-1",
                  "/usr/libexec/polkit-gnome-authentication-agent-1",
                  "/usr/libexec/polkit-mate-authentication-agent-1")

        for path in agents:
            if Path(path).exists():
                print(f"Launching for fallback session: {path}")
                self.polkit_agent_proc = subprocess.Popen([path])
                break

        # Start NM applet
        if shutil.which("nm-applet"):
            print(f"Launching for fallback session: nm-applet")
            self.nm_applet_proc = subprocess.Popen(["nm-applet"])

    def kill_fallback_helpers(self):
        if self.polkit_agent_proc is not None:
            print("Killing fallback polkit agent")
            self.polkit_agent_proc.terminate()
            try:
                self.polkit_agent_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.polkit_agent_proc.kill()
            self.polkit_agent_proc = None

        if self.nm_applet_proc is not None:
            print("Killing fallback nm-applet")
            self.nm_applet_proc.terminate()
            try:
                self.nm_applet_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.nm_applet_proc.kill()
            self.nm_applet_proc = None

    @async_function
    def monitor_memory(self):
        if psutil.pid_exists(self.cinnamon_pid):
            process = psutil.Process(self.cinnamon_pid)
            while True:
                delay = self.settings.get_int("check-frequency")
                limit = self.settings.get_int("memory-limit")
                time.sleep(delay)
                if not process.is_running():
                    break
                memory = ((process.memory_full_info().rss - process.memory_full_info().shared) / 1024 ** 2)
                if memory > limit:
                    print ("Memory threshold reached: %d" % memory)
                    self.log_kill(memory, limit, delay)
                    with self.kill_lock:
                        self.killed = True
                        process.kill()
                    break

    def restart_cinnamon(self):
        os.execvp(sys.argv[0], sys.argv)

    def log_kill(self, memory, limit, delay):
        try:
            directory = os.path.join(GLib.get_user_state_dir(), 'cinnamon')
        except AttributeError:
            directory = os.path.expanduser("~/.cinnamon")

        string = "%s - Usage: %d MB - Limit: %d MB - Freq: %s sec" % (time.strftime("%Y.%m.%d %H:%M:%S"), memory, limit, delay)
        if not os.path.exists(directory):
            os.mkdir(directory)
        with open(os.path.join(directory, "restarts.log"), "a") as log_file:
            print(string, file=log_file)

    @idle_function
    def confirm_restart(self):
        d = Gtk.MessageDialog(title=None, parent=None, flags=0, message_type=Gtk.MessageType.WARNING)
        d.add_buttons(_("No"),  Gtk.ResponseType.NO,
                      _("Yes"), Gtk.ResponseType.YES)
        d.set_keep_above(True)
        d.set_position(Gtk.WindowPosition.CENTER)
        d.set_title(_("Cinnamon just crashed"))
        box = d.get_content_area()
        checkbutton = Gtk.CheckButton(label=_("Disable downloaded applets, desklets and extensions"))
        d.set_markup("<span size='large'><b>%s</b></span>\n\n%s\n\n%s\n\n" %
                (
                    _("You are currently running in fallback mode."),
                    _("If you suspect the source of the crash is a local extension, applet or desklet, you can disable downloaded add-ons using the checkbox below."),
                    _("Do you want to restart Cinnamon?"),
                ))
        box.set_border_width(20)
        box.add(checkbutton)
        checkbutton.show_all()
        resp = d.run()
        d.destroy()
        if resp == Gtk.ResponseType.YES:
            if checkbutton.get_active():
                os.environ["CINNAMON_TROUBLESHOOT"] = "1"
            os.system("killall %s" % panel_process_name)
            self.kill_fallback_helpers()
            os.system("kill %d" % self.fb_pid)
            os.waitpid(self.fb_pid, 0)
            self.restart_cinnamon()
        else:
            Gtk.main_quit()

if __name__ == "__main__":
    setproctitle("cinnamon-launcher")
    Launcher()
