#!/usr/bin/python3

# -*- coding: utf-8 -*-
#
# This file is part of PRIME Indicator - indicator applet for NVIDIA Optimus laptops.
# Copyright (C) 2016 André Brait Carneiro Fabotti
#
# This work is based on the works of Alfred Neumayer and Clement Lefebvre.
#
# PRIME Indicator 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, or
# (at your option) any later version.
#
# PRIME Indicator 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 PRIME Indicator.  If not, see <http://www.gnu.org/licenses/>.

import sys
import os
import re
import signal
import subprocess
import configparser

# Uncomment the following import if you want to enable type hinting
# and you're running Python >3.5, or you installed the typing package
# using pip.
# from typing import List 

import gi.repository

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

DESKTOP = "kde"
APP_NAME = "PRIME Indicator"
HOME_DIR = os.getenv("HOME")  # type:str
CONFIG_DIR = "/etc/nvidia-prime"
LIB_PATH = "/usr/lib/prime-indicator/"
STD_ICONS_FOLDER = LIB_PATH + "icons/"
USER_ICONS_FOLDER = HOME_DIR + "/.config/prime-indicator/icons/"
CONFIG_PATH = HOME_DIR + "/.config/prime-indicator/prime-indicator.cfg"
NVIDIA_SETTINGS_PATH = "/usr/bin/nvidia-settings"

NVIDIA = "nvidia"
INTEL = "intel"

NVIDIA_COLOR = STD_ICONS_FOLDER + "prime-indicator-nvidia.svg"
INTEL_COLOR = STD_ICONS_FOLDER + "prime-indicator-intel.svg"
UNKNOWN_COLOR = STD_ICONS_FOLDER + "prime-indicator-unknown.svg"
NVIDIA_SYMBOLIC = USER_ICONS_FOLDER + "prime-indicator-nvidia-symbolic-themed.svg"
INTEL_SYMBOLIC = USER_ICONS_FOLDER + "prime-indicator-intel-symbolic-themed.svg"
UNKNOWN_SYMBOLIC = USER_ICONS_FOLDER + "prime-indicator-unknown-symbolic-themed.svg"


class Indicator:
    def __init__(self):

        self.config = configparser.ConfigParser()
        self.check_config_integrity()

        self.pm_enabled = (self.config.get("PowerManagement", "enabled").strip().lower() == "true")  # type:bool

        self.icon_path = LIB_PATH + "/icons"
        self.active_gpu = subprocess.getoutput("if grep -q 'prime-nvidia-config' /etc/X11/xorg.conf; then echo nvidia; else echo intel; fi")  # type:str
        self.theme_icons = self.config.get("Appearance", "iconset")  # type:str
        self.custom_color = None

        if self.active_gpu in [INTEL, NVIDIA]:
            if self.theme_icons == "theme-default":
                theme = Gtk.IconTheme.get_default()
                if theme.has_icon(self.active_gpu):
                    self.icon_name = theme.lookup_icon(self.active_gpu, max(theme.get_icon_sizes(self.active_gpu)),
                                                       0).get_filename()
                else:
                    self.icon_name = INTEL_COLOR if self.active_gpu == INTEL else NVIDIA_COLOR
                    print("ALERT: GTK Icon Theme does not provide an icon for " + (
                        "Intel" if self.active_gpu == INTEL else "NVIDIA") + ". Using default color icon instead!")
            elif self.theme_icons == "color":
                self.icon_name = INTEL_COLOR if self.active_gpu == INTEL else NVIDIA_COLOR
            else:
                if self.theme_icons.startswith("custom"):
                    self.custom_color = re.search("custom\((.*)\)", self.theme_icons).group(1)
                self.create_themed_icons()
                self.icon_name = INTEL_SYMBOLIC if self.active_gpu == INTEL else NVIDIA_SYMBOLIC
            self.icon_tooltip_text = "Active graphics card: " + ("Intel" if self.active_gpu == INTEL else "NVIDIA")
        else:
            if self.theme_icons == "theme-default":
                theme = Gtk.IconTheme.get_default()
                if theme.has_icon(self.active_gpu):
                    self.icon_name = theme.lookup_icon(self.active_gpu, max(theme.get_icon_sizes(self.active_gpu)),
                                                       0).get_filename()
                else:
                    self.icon_name = UNKNOWN_COLOR
                    print("ALERT: GTK Icon Theme does not provide an icon for Unknown. Using default color icon instead!")
            elif self.theme_icons == "color":
                self.icon_name = UNKNOWN_COLOR
            else:
                if self.theme_icons.startswith("custom"):
                    self.custom_color = re.search("custom\((.*)\)", self.theme_icons).group(1)
                self.create_themed_icons()
                self.icon_name = UNKNOWN_SYMBOLIC
            self.icon_tooltip_text = "Active graphics card: " + self.active_gpu

        self.menu = Gtk.Menu()
        item = Gtk.MenuItem(self.renderer_string())
        item.set_sensitive(False)
        item.show()
        self.menu.append(item)

        switch_in_use = Gtk.MenuItem("Quick switch graphics ...")
        switch_in_use.connect("activate", self.switch)
        self.menu.append(switch_in_use)

        separator_in_use = Gtk.SeparatorMenuItem()
        self.menu.append(separator_in_use)

        self.toggle_power_management_enable = Gtk.CheckMenuItem("Enable NVIDIA GPU Power Management")
        self.toggle_power_management_enable.set_active(self.pm_enabled)
        self.toggle_power_management_enable.connect("toggled", self.toggle_pm)
        self.menu.append(self.toggle_power_management_enable)

        self.info_power_management = Gtk.MenuItem()
        self.info_power_management.set_sensitive(False)
        self.menu.append(self.info_power_management)


        self.set_nv_pm_labels()
        if self.active_gpu in [INTEL, NVIDIA]:
            switch_in_use.show()
            separator_in_use.show()
            if self.active_gpu == INTEL:
                self.toggle_power_management_enable.show()
                if self.pm_enabled:
                    self.info_power_management.show()

        item = Gtk.SeparatorMenuItem()
        item.show()
        self.menu.append(item)

        item = Gtk.MenuItem("Open NVIDIA Settings")
        item.connect("activate", self.run_nvidia_settings)
        item.show()
        self.menu.append(item)

        item = Gtk.SeparatorMenuItem()
        item.show()
        self.menu.append(item)

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

        self.icon = AppIndicator3.Indicator.new(APP_NAME, "",
                                                AppIndicator3.IndicatorCategory.APPLICATION_STATUS)
        self.icon.set_status(AppIndicator3.IndicatorStatus.ACTIVE)
        self.icon.set_icon(self.icon_name)
        self.icon.set_title(self.icon_tooltip_text)
        self.icon.set_menu(self.menu)

    def create_themed_icons(self):
        window = Gtk.Window()  # type:Gtk.Window
        fg_color_str = "#bebebe"
        if self.custom_color:
            fg_color_str = self.custom_color
        else:
            style_context = window.get_style_context()
            style_fg_color = style_context.lookup_color('fg_color')  # type:tuple
            if style_fg_color[0]:
                fg_color_str = "#{r:02x}{g:02x}{b:02x}".format(r=int(style_fg_color[1].red * 255),
                                                               g=int(style_fg_color[1].green * 255),
                                                               b=int(style_fg_color[1].blue * 255))

        os.makedirs(USER_ICONS_FOLDER, exist_ok=True)
        with open(STD_ICONS_FOLDER + "prime-indicator-nvidia-symbolic.svg", "r") as fin:
            with open(NVIDIA_SYMBOLIC, "w") as fout:
                for line in fin:
                    fout.write(line.replace("#bebebe", fg_color_str))
        with open(STD_ICONS_FOLDER + "prime-indicator-intel-symbolic.svg", "r") as fin:
            with open(INTEL_SYMBOLIC, "w") as fout:
                for line in fin:
                    fout.write(line.replace("#bebebe", fg_color_str))

    def write_default_config(self):
        os.makedirs(os.path.dirname(CONFIG_PATH), exist_ok=True)
        self.config.clear()
        self.config.add_section("PowerManagement")
        self.config.add_section("Appearance")
        self.config.set("PowerManagement", "enabled", "true")
        self.config.set("Appearance", "iconset", "theme-default")
        with open(CONFIG_PATH, "w") as configfile:
            self.config.write(configfile)

    def check_config_integrity(self):
        if not os.path.exists(CONFIG_PATH) or not os.path.isfile(CONFIG_PATH):
            self.write_default_config()
        else:
            try:
                self.config.read(CONFIG_PATH)
                iconset = self.config.get("Appearance", "iconset")
                if iconset not in ["symbolic", "color", "theme-default"] and \
                        not re.search('^custom\(#([a-fA-F0-9]{2}){3}\)$', iconset) or \
                        self.config.get("PowerManagement", "enabled") not in ["true", "false"]:
                    self.write_default_config()
            except (configparser.NoSectionError, configparser.NoOptionError):
                self.write_default_config()

    def toggle_pm(self, dude):
        self.pm_enabled = self.toggle_power_management_enable.get_active()
        self.config.set(
            "PowerManagement", "enabled", str(self.pm_enabled))
        with open(CONFIG_PATH, "w") as configfile:
            self.config.write(configfile)
        if self.pm_enabled and self.active_gpu == INTEL:
            self.info_power_management.show()
            self.switch_power_management.show()
        else:
            self.info_power_management.hide()
            self.switch_power_management.hide()

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

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

    def execute(self):
        Gtk.main()

    def logout(self):
        user=os.environ['SUDO_USER']
        print("logging out")
        print(os.environ['SUDO_USER'])
        if DESKTOP.startswith("xfce"):
            os.system("xfce4-session-logout --logout")
            restart()
        elif DESKTOP.startswith("kde"):
            os.system("sudo -u " + user+ " qdbus org.kde.ksmserver /KSMServer logout 0 0 0")
        elif DESKTOP.startswith("lxde"):
            os.system("lxsession-logout --prompt 'Please click the Log Out button to continue'")
        elif DESKTOP.startswith("x-cinnamon"):
            os.system("cinnamon-session-quit --logout --no-prompt")
        elif DESKTOP.startswith("mate"):
            os.system("mate-session-save --logout")
        elif DESKTOP.startswith("budgie"):
            os.system("budgie-session --logout")
        elif DESKTOP.startswith("lxqt"):
            os.system("lxqt-leave --logout ")
        elif DESKTOP.startswith("gnome") or env.startswith("pantheon") or env.startswith("unity"):
            os.system("gnome-session-quit --logout --no-prompt")
        elif DESKTOP.startswith("i3"):
            os.system("i3-msg exit")
        else:
            message = "It seems you're running an unsupported Desktop Environment. " + \
                      "Please manually log out and then log in again to complete the switch."
            dialog = Gtk.MessageDialog(None, Gtk.DIALOG_MODAL, Gtk.MESSAGE_ERROR, Gtk.BUTTONS_OK, message)
            dialog.set_deletable(False)
            dialog.run()
            restart()

    def renderer_string(self) -> str:
        out = subprocess.getoutput("glxinfo | grep 'OpenGL renderer string'")  # type:str
        out = out.replace("OpenGL renderer string", "Using")
        return out

    def nv_power_string(self) -> str:
        return "NVIDIA GPU is powered " + ("ON" if self.is_nvidia_on() else "OFF")

    def nv_power_switch_string(self) -> str:
        return "Force NVIDIA GPU to " + ("power OFF" if self.is_nvidia_on() else "power ON")

    def is_nvidia_on(self) -> bool:
        out = subprocess.getoutput("cat /proc/acpi/bbswitch")
        return out.lower().endswith("on")

    def switch_gpu(self) -> None:
        if self.active_gpu in [INTEL, NVIDIA]:
            print("Switching from: " + self.active_gpu)
            if self.active_gpu == INTEL:
               os.system( "sudo ln -sf " + CONFIG_DIR + "/xorg.nvidia.conf /etc/X11/xorg.conf" )
            else:
               os.system( "sudo ln -sf " + CONFIG_DIR + "/xorg.intel.conf /etc/X11/xorg.conf" )


    def switch_nv_power(self, dude) -> None:
        if self.is_nvidia_on():
            self.turn_nv_off()
        else:
            self.turn_nv_on()

    def turn_nv_on(self) -> None:
        os.system(SCRIPT_CMD + " nvidia on")
        self.set_nv_pm_labels()

    def turn_nv_off(self) -> None:
        os.system(SCRIPT_CMD + " nvidia off")
        self.set_nv_pm_labels()

    def set_nv_pm_labels(self) -> None:
        self.info_power_management.set_label(self.nv_power_string())

    def ignore(self, *args):
        return Gtk.ResponseType.CANCEL

    def show_reboot_dialog(self):
        msg_nvidia = "dedicated NVIDIA GPU"
        msg_intel = "integrated Intel GPU"
        message = "You need to log out to switch from the " + \
                  (msg_intel if (self.active_gpu == INTEL) else msg_nvidia) + \
                  " to the " + (msg_nvidia if (self.active_gpu == INTEL) else msg_intel) + \
                  ". Save your work before clicking the Log Out button below."
        dialog = Gtk.MessageDialog(None, Gtk.DialogFlags.MODAL, Gtk.MessageType.INFO,
                                   Gtk.ButtonsType.NONE, message)
        dialog.set_deletable(False)
        dialog.connect("delete_event", self.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 switch(self, dude) -> None:
        response = self.show_reboot_dialog()
        if response != Gtk.ResponseType.CANCEL:
            self.switch_gpu()
            self.logout()


def kill_other_instances() -> None:
    otherpid = subprocess.getoutput("pgrep -f prime-indicator")  # type:bytes
    if otherpid:
        otherpid = str(otherpid)  # type:str
        pidlist = otherpid.splitlines()  # type:List[str]
        for pid in pidlist:
            if pid and pid.isnumeric():
                pid = int(pid)  # type:int
                if pid != os.getpid():
                    try:
                        os.kill(pid, signal.SIGTERM)
                        os.kill(pid, signal.SIGKILL)
                    except ProcessLookupError:
                        pass


def run() -> None:
    Indicator().execute()


if __name__ == "__main__":

    # If nvidia-prime isn't installed or isn't supported, exit cleanly
#    if not os.path.exists(NVIDIA_SETTINGS_PATH) or not os.path.exists(PRIME_SELECT_PATH):
#        print("bla");
#        sys.exit(0)

#    if subprocess.getoutput("prime-supported 2>/dev/null") != "yes":
#        sys.exit(0)

    kill_other_instances()
    run()
