#!/usr/bin/env python3

import gettext
import gi
import os
import setproctitle
import subprocess
import sys
from subprocess import DEVNULL, PIPE, STDOUT

gi.require_version('Gdk', '3.0')
from gi.repository import Gdk, GdkPixbuf

gi.require_version('Gtk', '3.0')
from gi.repository import Gtk

gi.require_version('Notify', '0.7')
from gi.repository import Notify

gi.require_version('AppIndicator3', '0.1')
from gi.repository import AppIndicator3

__VERSION__ = '19.10.4'

# i18n
gettext.install("mate-optimus", "/usr/share/locale")
skip_prime_support = False

def get_active_gpu():
    process = subprocess.Popen(['prime-select', 'query'], stdout=PIPE)
    out, err = process.communicate()
    return out.decode('utf-8')

def has_nvidia_settings():
    # If NVIDIA Prime is not installed or isn't supported then exit.
    if (os.path.exists("/usr/bin/nvidia-settings") and os.path.exists("/usr/bin/prime-select")):
        print('  - nvidia-settings and prime-select detected.')
        return True
    else:
        print('  - nvidia-settings and prime-select not detected.')
        return False

def prime_capable():
    if skip_prime_support:
        print('  - prime-supported: Skipping.')
        return True

    prime_supported = subprocess.Popen(['prime-supported'], stdout=PIPE, stderr=DEVNULL)
    out, err = prime_supported.communicate()
    if 'yes' in out.decode('utf-8'):
        print('  - prime-supported: Yes')
        return True
    else:
        print('  - prime-supported: No')
        return False

def nvidia_detected(capability):
    if capability == 'nvidia' or capability == 'on-demand':
        # If we're check for nvidia let's make real sure we have an nvidia GPU.
        nvidia_detector = subprocess.Popen(['nvidia-detector'], stdout=PIPE, stderr=DEVNULL)
        out, err = nvidia_detector.communicate()
        driver = out.decode('utf-8').lower().strip()
        if driver == 'none':
            print('  - nvidia-detector: None')
            return False
        else:
            print('  - nvidia-detector: ' + driver)
            return True
    else:
        print('  - nvidia-detector: Skipped')
        return True

def has_gpu_screen(capability):
    if capability == 'on-demand':
        providers = subprocess.Popen(['xrandr', '--listproviders'], stdout=PIPE, stderr=DEVNULL)
        for line in providers.stdout.readlines():
            text = line.rstrip().decode('utf-8')
            if 'name:NVIDIA-G' in text:
                print('  - xrandr: NVIDIA GPU screen detected.')
                return True
        print('  - xrandr: NVIDIA GPU screen not detected.')
        return False
    else:
        print('  - xrandr: NVIDIA GPU screen check skipped.')
        return True

def gpu_loaded(gpu):
    # Does prime-switch report that the GPU is available?
    if gpu == 'intel' or gpu == 'amdgpu' or gpu == 'radeon':
        gpu_loaded = gpu + ' loaded? yes'
        prime_switch = subprocess.Popen(['prime-switch', '-help'], stdout=PIPE, stderr=DEVNULL)
        for line in prime_switch.stdout.readlines():
            text = line.rstrip().decode('utf-8')
            if gpu_loaded in text:
                print('  - ' + gpu + ' loaded: Yes')
                return True
        print('  - ' + gpu + ' loaded: No')
        return False
    else:
        print('  - ' + gpu + ' loaded: n/a')
        return True

def prime_select_features(capability):
    # prime-select outputs to stderr
    prime_select = subprocess.Popen(['prime-select'], stdout=PIPE, stderr=PIPE)
    out, err = prime_select.communicate()
    if capability in err.decode('utf-8'):
        print('  - prime-select: ' + capability + ' detected.')
        return True
    else:
        print('  - prime-select: ' + capability + ' not detected.')
        return False

def check_capability(capability):
    print('Checking: ' + capability)
    # Some other things to check for a compliant system here:
    #  * https://download.nvidia.com/XFree86/Linux-x86_64/435.21/README/primerenderoffload.html

    if gpu_loaded(capability) and \
       prime_capable() and \
       nvidia_detected(capability) and \
       prime_select_features(capability) and \
       has_gpu_screen(capability):
        print('  - ' + capability + ': is supported')
        return True
    else:
        print('  - ' + capability + ': is not supported')
        return False

def confirm_ignore(*args):
    return Gtk.ResponseType.CANCEL

def confirm_gpu_switch(gpu):
    message = _('Do you want to switch GPU now?') + ' ' + _('You need to log out and then log back in to switch the GPU to') + ' ' + gpu

    dialog = Gtk.MessageDialog(None, Gtk.DialogFlags.MODAL, Gtk.MessageType.INFO, Gtk.ButtonsType.NONE, message)
    dialog.set_deletable(False)
    dialog.connect("delete_event", confirm_ignore)
    dialog.add_button(_("Log Out"), Gtk.ResponseType.OK)
    dialog.add_button(Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL)
    response = dialog.run()
    dialog.destroy()
    return response

def run_nvidia_settings(arg=None):
    subprocess.Popen(["nvidia-settings", "-page", "PRIME Profiles"])

def session_logout():
    current_desktop = os.environ.get("XDG_CURRENT_DESKTOP").lower().strip()
    if current_desktop.startswith('mate'):
        subprocess.Popen(["mate-session-save", "--logout-dialog", "--gui"])
    elif current_desktop.startswith('budgie'):
        subprocess.Popen(["gnome-session-quit", "--logout", "--no-prompt"])
    elif current_desktop.startswith('gnome'):
         subprocess.Popen(["gnome-session-quit", "--logout"])
    elif current_desktop.startswith('kde'):
        subprocess.Popen(["qdbus", "org.kde.ksmserver", "/KSMServer", "logout", "0", "0", "0"])
    elif current_desktop.startswith('lxqt'):
        subprocess.Popen(["lxqt-leave", "--logout"])
    elif current_desktop.startswith('x-cinnamon'):
        subprocess.Popen(["cinnamon-session-quit", "--logout"])
    elif current_desktop.startswith('xfce'):
        subprocess.Popen(["xfce4-session-logout", "--logout"])
    else:
        pass

def switch_to_intel(arg=None):
    response = confirm_gpu_switch(_("Intel (Power Saving)"))
    if response != Gtk.ResponseType.CANCEL:
        subprocess.call(["pkexec", "prime-select", "intel"])
        session_logout()

def switch_to_nvidia(arg=None):
    response = confirm_gpu_switch(_("NVIDIA (Performance Mode)"))
    if response != Gtk.ResponseType.CANCEL:
        subprocess.call(["pkexec", "prime-select", "nvidia"])
        session_logout()

def switch_to_ondemand(arg=None):
    response = confirm_gpu_switch(_("NVIDIA (On-Demand)"))
    if response != Gtk.ResponseType.CANCEL:
        subprocess.call(["pkexec", "prime-select", "on-demand"])
        session_logout()

class Indicator:
    def __init__(self):
        self.intel_available = check_capability('intel')
        self.nvidia_available = check_capability('nvidia')
        self.ondemand_available = check_capability('on-demand')
        self.active_gpu = get_active_gpu()

        if (self.intel_available and 'intel' in self.active_gpu):
            current_icon = "optimus-indicator-intel"
            current_gpu = _("Intel (Power Saving)")
        elif (self.nvidia_available and 'nvidia' in self.active_gpu):
            current_icon = "optimus-indicator-nvidia"
            current_gpu = _("NVIDIA (Performance Mode)")
        elif (self.ondemand_available and 'on-demand' in self.active_gpu):
            current_icon = "optimus-indicator-nvidia"
            current_gpu = _("NVIDIA (On-Demand)")
        else:
            current_icon = "optimus-indicator-unknown"
            current_gpu = _("Unknown")

        # If we can derive the OpenGL renderer use that for current_gpu
        visualinfo = subprocess.Popen('visualinfo', stdout=PIPE, stderr=STDOUT)
        for line in visualinfo.stdout.readlines():
            text = line.rstrip().decode('utf-8')
            if 'OpenGL renderer string:' in text:
                current_gpu = text.replace('OpenGL renderer string:','').strip()
                break

        menu = Gtk.Menu()

        self.icon = AppIndicator3.Indicator.new('mate-optimus', current_icon, AppIndicator3.IndicatorCategory.HARDWARE)
        self.icon.set_status(AppIndicator3.IndicatorStatus.ACTIVE)

        self.icon.set_title(current_gpu)
        active = Gtk.MenuItem(label=current_gpu)
        active.set_sensitive(False)
        menu.append(active)

        menu.append(Gtk.SeparatorMenuItem())

        if self.intel_available:
            intel = Gtk.MenuItem(label=_("Switch to: ") + _("Intel (Power Saving)"))
            intel.connect("activate", switch_to_intel)
            if 'intel' in self.active_gpu:
                intel.set_sensitive(False)
            menu.append(intel)

        if self.nvidia_available:
            nvidia = Gtk.MenuItem(label=_("Switch to: ") + _("NVIDIA (Performance Mode)"))
            nvidia.connect("activate", switch_to_nvidia)
            if 'nvidia' in self.active_gpu:
                nvidia.set_sensitive(False)
            menu.append(nvidia)

        if self.ondemand_available:
            ondemand = Gtk.MenuItem(label=_("Switch to: ") + _("NVIDIA (On-Demand)"))
            ondemand.connect("activate", switch_to_ondemand)
            if 'on-demand' in self.active_gpu:
                ondemand.set_sensitive(False)
            menu.append(ondemand)

        if self.intel_available or self.nvidia_available or self.ondemand_available:
            menu.append(Gtk.SeparatorMenuItem())

        if has_nvidia_settings():
            settings = Gtk.MenuItem(label=_("NVIDIA Settings"))
            settings.connect("activate", run_nvidia_settings)
            menu.append(settings)
            menu.append(Gtk.SeparatorMenuItem())

        about_app = Gtk.MenuItem(label=_("About"))
        about_app.connect("activate", self.show_about)
        menu.append(about_app)

        quit_app = Gtk.MenuItem(label=_("Quit"))
        quit_app.connect("activate", self.terminate)
        menu.append(quit_app)
        self.icon.set_menu(menu)
        menu.show_all()

    def show_about(self, widget):
        about = Gtk.AboutDialog()
        about.set_program_name("MATE Optimus")
        about.set_version(__VERSION__)
        about.set_authors(["Martin Wimpress", "Clement Lefebvre"])
        about.set_copyright('© 2015-2019 Ubuntu MATE Team')
        about.set_comments("NVIDIA Optimus GPU switcher.")
        about.set_website("https://github.com/ubuntu-mate/mate-optimus")
        about.set_website_label("https://github.com/ubuntu-mate/mate-optimus")
        about.set_license_type(Gtk.License.GPL_3_0)
        try:
            about.set_logo(GdkPixbuf.Pixbuf.new_from_file_at_size(os.path.join('usr', 'share', 'pixmaps', 'optimus-indicator-nvidia.svg'), 64, 64))
        except:
            about.set_logo_icon_name('optimus-indicator-nvidia')

        about.run()
        about.destroy()

    def terminate(self, window = None, data = None):
        Gtk.main_quit()

def offload_launch(mode='', command=['']):
    # Make sure mode is sane.
    if mode != 'offload-glx' and mode != 'offload-vulkan':
        mode = 'offload-none'

    # Validate the command we've be passed.
    if isinstance(command, list):
        command_str = ' '.join(command)
    else:
        command_str = ''

    if not len(command_str):
        print(mode + ': No command given, doing nothing.')
        return 1

    # Copy the environment so we don't pollute it.
    offload_env = os.environ.copy()

    if mode == 'offload-glx':
        offload_env["__NV_PRIME_RENDER_OFFLOAD"] = "1"
        offload_env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
        print(mode + ': Launching "' + command_str + '" with NVIDIA offload.')
    elif mode == 'offload-vulkan':
        offload_env["__NV_PRIME_RENDER_OFFLOAD"] = "1"
        print(mode + ': Launching "' + command_str + '" with NVIDIA offload.')
    else:
        print(mode + ': Launching "' + command_str +'" with no offload.')

    # Spawn the process in the users default shell using the modified environment
    offload = subprocess.Popen(command_str, env=offload_env, shell=True, executable=os.getenv('SHELL', '/bin/sh'))
    offload.wait()
    return offload.returncode

if __name__ == "__main__":
    executable = os.path.basename(__file__)
    setproctitle.setproctitle(executable)
    print(executable + ': launched.')

    if (has_nvidia_settings() and prime_capable()) or skip_prime_support:
        print(executable + ': NVIDIA Optimus supported.')
        if executable == 'mate-optimus-applet':
            Indicator()
            Gtk.main()
        elif executable == 'offload-glx' or executable == 'offload-vulkan':
            if ('on-demand' in get_active_gpu()):
                if check_capability('on-demand'):
                    if executable == 'offload-glx':
                        print(executable + ': NVIDIA On-Demand is enabled.')
                        returncode = offload_launch(executable, sys.argv[1:])
                        sys.exit(returncode)
                    elif executable == 'offload-vulkan':
                        print(executable + ': NVIDIA On-Demand is enabled.')
                        returncode = offload_launch(executable, sys.argv[1:])
                        sys.exit(returncode)
                else:
                    print(executable + ": NVIDIA On-Demand is enabled but the system doesn't appear to be compatible.")
                    returncode = offload_launch(None, sys.argv[1:])
                    sys.exit(returncode)
            else:
                print(executable + ': NVIDIA On-Demand is not currently enabled.')
                returncode = offload_launch(None, sys.argv[1:])
                sys.exit(returncode)
        else:
            print(executable + ': I have an identity crisis. Send help!')
    else:
        print(executable + ': NVIDIA Optimus is not supported.')
