#!/usr/bin/python3
import gettext
import gi
import os
import subprocess
import sys
from packaging import version

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

# i18n
gettext.install("nvidia-prime-applet", "/usr/share/locale")

NVIDIA_MODE = _("NVIDIA (Performance Mode)")
ON_DEMAND_MODE = _("NVIDIA On-Demand")
INTEL_MODE = _("Intel (Power Saving Mode)")
AMD_MODE = _("AMD (Power Saving Mode)")

MIN_ON_DEMAND_VERSION = "435.17"

INTEL_PCI_ID = "8086"
NVIDIA_PCI_ID = "10de"
AMD_PCI_ID = "1002"

def get_output(commands):
    process = subprocess.Popen(commands, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    out, err = process.communicate()
    return out.decode('utf-8').strip()

class Tray:
    def __init__(self, integrated_is_intel):

        self.icon = XApp.StatusIcon()
        self.icon.set_name("nvidia-prime")

        hybrid_supported = self.should_show_hybrid()

        # Find GPU name
        renderer = None
        try:
            renderer = subprocess.check_output("glxinfo | grep -i 'OpenGL renderer'", shell=True).decode("UTF-8").strip().split(": ")[1]
        except:
            pass

        if integrated_is_intel:
            integrated_icon_name = "prime-tray-intel-symbolic"
            integrated_mode = INTEL_MODE
        else:
            integrated_icon_name = "prime-tray-amd-symbolic"
            integrated_mode = AMD_MODE

        # Find active mode
        nvidia = Gtk.MenuItem(label=_("Switch to: %s") % NVIDIA_MODE)
        nvidia.connect("activate", self.switch, 'nvidia', NVIDIA_MODE)

        if hybrid_supported:
            ondemand = Gtk.MenuItem(label=_("Switch to: %s") % ON_DEMAND_MODE)
            ondemand.connect("activate", self.switch, 'on-demand', ON_DEMAND_MODE)

        integrated = Gtk.MenuItem(label=_("Switch to: %s") % integrated_mode)
        integrated.connect("activate", self.switch, 'intel', integrated_mode)

        active_gpu = get_output(["prime-select", "query"])
        if (active_gpu == "nvidia"):
            self.icon.set_icon_name("prime-tray-nvidia-symbolic")
            mode = NVIDIA_MODE
        elif (active_gpu == "on-demand"):
            self.icon.set_icon_name(integrated_icon_name)
            mode = ON_DEMAND_MODE
        elif (active_gpu == "intel"):
            self.icon.set_icon_name(integrated_icon_name)
            mode = integrated_mode
        else:
            self.icon.set_icon_name("dialog-error-symbolic")
            mode = _("Unknown mode")

        menu = Gtk.Menu()

        if renderer is None:
            self.icon.set_tooltip_text(mode)
        else:
            self.icon.set_tooltip_text("%s\n%s" % (renderer, mode))
            item = Gtk.MenuItem(label=renderer)
            item.set_sensitive(False)
            menu.append(item)
            menu.append(Gtk.SeparatorMenuItem())

        item = Gtk.MenuItem(label=_("Active profile: %s") % mode)
        item.set_sensitive(False)
        menu.append(item)

        if mode != integrated_mode:
            menu.append(integrated)
        if mode != NVIDIA_MODE:
            menu.append(nvidia)
        if mode != ON_DEMAND_MODE and hybrid_supported:
            menu.append(ondemand)

        menu.append(Gtk.SeparatorMenuItem())

        item = Gtk.MenuItem(label=_("NVIDIA Settings"))
        item.connect("activate", self.run_nvidia_settings)
        menu.append(item)
        menu.append(Gtk.SeparatorMenuItem())

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

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

        self.icon.set_primary_menu(menu)
        self.icon.set_secondary_menu(menu)

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

    def should_show_hybrid(self):
        # > dkms status
        # nvidia, 440.100, 5.4.0-40-generic, x86_64: installed
        try:
            for line in get_output(["dkms", "status"]).splitlines():
                if "nvidia" in line:
                    min_version = version.parse(MIN_ON_DEMAND_VERSION)

                    current_version = version.parse(line.split(",")[1].strip())
                    print("Required nvidia driver version for on-demand: %s" % str(min_version))
                    print("Detected nvidia driver version: %s" % str(current_version))
                    if current_version < min_version:
                        print("On-demand not available")
                        return False
                    break
        except:
            pass

        return True

    def dialog_closed(self, widget, event):
        return Gtk.ResponseType.NO

    def switch(self, widget, mode, mode_description):
        dialog = Gtk.MessageDialog(parent=None, message_type=Gtk.MessageType.INFO, buttons=Gtk.ButtonsType.YES_NO, text=_("Are you sure you want to switch to %s?") % mode_description)
        dialog.format_secondary_text(_('Changes will take effect after you log out and log back in.'))
        dialog.set_deletable(False)
        dialog.set_title("NVIDIA Optimus")
        dialog.set_icon_name("prime-tray-nvidia")
        dialog.set_skip_taskbar_hint(False)
        dialog.set_skip_pager_hint(False)
        dialog.connect("delete_event", self.dialog_closed)
        response = dialog.run()
        dialog.destroy()
        if response == Gtk.ResponseType.YES:
            subprocess.call(["pkexec", "prime-select", mode])

    def about(self, widget):
        about = Gtk.AboutDialog()
        about.set_program_name("NVIDIA Optimus")
        about.set_website("https://github.com/linuxmint/nvidia-prime-applet")
        about.set_website_label("https://github.com/linuxmint/nvidia-prime-applet")
        about.set_license_type(Gtk.License.GPL_3_0)
        about.set_logo_icon_name('prime-tray-nvidia')

        about.run()
        about.destroy()

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

if __name__ == "__main__":

    # If nvidia-prime isn't installed or isn't supported, exit cleanly
    if not (os.path.exists("/usr/bin/nvidia-settings") and os.path.exists("/usr/bin/prime-select")):
        print("Aborting: Check that the nvidia-settings and nvidia-prime packages are installed.")
        sys.exit(0)

    has_intel = False
    has_nvidia = False
    has_amd = False
    path = '/var/lib/ubuntu-drivers-common/last_gfx_boot'
    if os.path.isfile(path):
        with open(path, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith("%s:" % INTEL_PCI_ID):
                    has_intel = True
                    print("Detected Intel GPU:", line)
                elif line.startswith("%s:" % NVIDIA_PCI_ID):
                    has_nvidia = True
                    print("Detected NVIDIA GPU:", line)
                elif line.startswith("%s:" % AMD_PCI_ID):
                    has_amd = True
                    print("Detected AMD GPU:", line)
    else:
        print("%s not found!" % path)
        sys.exit(0)

    if not has_nvidia:
        print("No NVIDIA card detected!")
        sys.exit(0)

    if has_intel or has_amd:
        Tray(integrated_is_intel=has_intel)
        Gtk.main()
    else:
        print("No Intel/AMD card detected!")
        sys.exit(0)
