#!/usr/bin/python3
""" amdgpu-monitor  -  Displays current status of all active GPUs

    A utility to give the current state of all compatible AMD GPUs. The default behavior
    is to continuously update a text based table in the current window until Ctrl-C is
    pressed.  With the *--gui* option, a table of relevant parameters will be updated
    in a Gtk window.  You can specify the delay between updates with the *--sleep N*
    option where N is an integer > zero that specifies the number of seconds to sleep
    between updates.  The *--no_fan* option can be used to disable the reading and display
    of fan information.  The *--log* option is used to write all monitor data to a psv log
    file.  When writing to a log file, the utility will indicate this in red at the top of
    the window with a message that includes the log file name. The *--plot* will display a
    plot of critical GPU parameters which updates at the specified *--sleep N* interval. If
    you need both the plot and monitor displays, then using the --plot option is preferred
    over running both tools as a single read of the GPUs is used to update both displays.
    The *--ltz* option results in the use of local time instead of UTC.

    Copyright (C) 2019  RueiKe

    This program 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.

    This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
"""
__author__ = 'RueiKe'
__copyright__ = 'Copyright (C) 2019 RueiKe'
__credits__ = ['Craig Echt - Testing, Debug, Verification, and Documentation']
__license__ = 'GNU General Public License'
__program_name__ = 'amdgpu-monitor'
__version__ = 'v3.0.0'
__maintainer__ = 'RueiKe'
__status__ = 'Stable Release'
__docformat__ = 'reStructuredText'
# pylint: disable=multiple-statements
# pylint: disable=line-too-long

import argparse
import subprocess
import threading
import os
import sys
import shlex
import time
import signal

try:
    import gi
except ModuleNotFoundError as error:
    print('gi import error: {}'.format(error))
    print('gi is required for {}'.format(__program_name__))
    print('   In a venv, first install vext:  pip install --no-cache-dir vext')
    print('   Then install vext.gi:  pip install --no-cache-dir vext.gi')
    sys.exit(0)
gi.require_version('Gtk', '3.0')
from gi.repository import GLib, Gtk, Gdk

from GPUmodules import GPUmodule as gpu
from GPUmodules import env


def ctrl_c_handler(target_signal, frame):
    """
    Signal catcher for ctrl-c to exit monitor loop.
    :param target_signal:
    :type target_signal: Signals
    :param frame:
    :return: None
    """
    if env.GUT_CONST.DEBUG:
        print('ctrl_c_handler (ID: {}) has been caught. Setting quit flag...'.format(target_signal))
    else:
        print('Setting quit flag...')
    MonitorWindow.quit = True


signal.signal(signal.SIGINT, ctrl_c_handler)

# SEMAPHORE ############
UD_SEM = threading.Semaphore()
########################


class MonitorWindow(Gtk.Window):
    """
    Custom PAC Gtk window.
    """
    quit = False

    def __init__(self, gpu_list, devices):

        Gtk.Window.__init__(self, title='amdgpu-monitor')
        self.set_border_width(1)
        icon_file = os.path.join(env.GUT_CONST.icon_path, 'amdgpu-monitor.icon.png')
        if env.GUT_CONST.DEBUG: print('Icon file: [{}]'.format(icon_file))
        if os.path.isfile(icon_file):
            self.set_icon_from_file(icon_file)
        grid = Gtk.Grid()
        grid.override_background_color(Gtk.StateType.NORMAL, Gdk.RGBA(1, 1, 1, 1))
        self.add(grid)

        col = 0
        row = 0
        num_amd_gpus = gpu_list.num_gpus()['total']
        if env.GUT_CONST.LOG:
            log_label = Gtk.Label()
            log_label.set_markup('<big><b> Logging to:    </b>{}</big>'.format(env.GUT_CONST.log_file))
            log_label.override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(1.0, 1.0, 1.0, 1.0))
            lbox = Gtk.Box(spacing=6)
            lbox.override_background_color(Gtk.StateType.NORMAL, Gdk.RGBA(.60, .20, .20, 1.0))
            lbox.set_property('margin-top', 1)
            lbox.set_property('margin-bottom', 1)
            lbox.set_property('margin-right', 1)
            lbox.set_property('margin-left', 1)
            lbox.pack_start(log_label, True, True, 0)
            grid.attach(lbox, 0, row, num_amd_gpus+1, 1)
        row += 1
        row_start = row

        row = row_start
        row_labels = {'card_num': Gtk.Label()}
        row_labels['card_num'].set_markup('<b>Card #</b>')
        row_labels['card_num'].override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(1.0, 1.0, 1.0, 1.0))
        for k, v in gpu_list.table_param_labels().items():
            row_labels[k] = Gtk.Label()
            row_labels[k].set_markup('<b>{}</b>'.format(v))
            row_labels[k].override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(1.0, 1.0, 1.0, 1.0))
        for k, v in row_labels.items():
            lbox = Gtk.Box(spacing=6)
            lbox.override_background_color(Gtk.StateType.NORMAL, Gdk.RGBA(.20, .40, .60, 1.0))
            lbox.set_property('margin-top', 1)
            lbox.set_property('margin-bottom', 1)
            lbox.set_property('margin-right', 1)
            lbox.set_property('margin-left', 1)
            v.set_property('margin-top', 1)
            v.set_property('margin-bottom', 1)
            v.set_property('margin-right', 4)
            v.set_property('margin-left', 4)
            lbox.pack_start(v, True, True, 0)
            grid.attach(lbox, col, row, 1, 1)
            v.set_alignment(0, 0.5)
            row += 1
        for v in gpu_list.list.values():
            devices[v.prm.uuid] = {'card_num':  Gtk.Label(label='card{}'.format(v.get_params_value('card_num')))}
            for cv in gpu_list.table_param_labels():
                devices[v.prm.uuid][cv] = Gtk.Label(label=v.get_params_value(str(cv)))
                devices[v.prm.uuid][cv].set_width_chars(10)

        for dv in devices.values():
            col += 1
            row = row_start
            for lv in dv.values():
                lv.set_text('')
                lbox = Gtk.Box(spacing=6)
                lbox.override_background_color(Gtk.StateType.NORMAL, Gdk.RGBA(.06, .06, .06, .06))
                lbox.set_property('margin-top', 1)
                lbox.set_property('margin-bottom', 1)
                lbox.set_property('margin-right', 1)
                lbox.set_property('margin-left', 1)
                lv.set_property('margin-top', 1)
                lv.set_property('margin-bottom', 1)
                lv.set_property('margin-right', 3)
                lv.set_property('margin-left', 3)
                lv.set_width_chars(17)
                lbox.pack_start(lv, True, True, 0)
                grid.attach(lbox, col, row, 1, 1)
                row += 1

    def set_quit(self, _arg2, _arg3):
        """
        Set quit flag when Gtk quit is selected.
        :param _arg2:
        :param _arg3:
        :return: None
        """
        self.quit = True


def update_data(gpu_list, devices, cmd):
    """
    Update monitor data with data read from GPUs.
    :param gpu_list: A gpuList object with all gpuItems
    :type gpu_list: gpuList
    :param devices: A dictionary linking Gui items with data.
    :type devices: dict
    :param cmd: Subprocess return from running plot.
    :type cmd: subprocess.Popen
    :return: None
    """
    # SEMAPHORE ############
    if not UD_SEM.acquire(blocking=False):
        if env.GUT_CONST.DEBUG: print('Update while updating, skipping new update')
        return
    ########################
    gpu_list.read_gpu_sensor_data(data_type='DynamicM')
    gpu_list.read_gpu_sensor_data(data_type='StateM')
    if env.GUT_CONST.LOG:
        gpu_list.print_log(env.GUT_CONST.log_file_ptr)
    if env.GUT_CONST.PLOT:
        try:
            gpu_list.print_plot(cmd.stdin)
        except (OSError, KeyboardInterrupt) as except_err:
            if env.GUT_CONST.DEBUG:
                print('amdgpu-plot has closed: [{}]'.format(except_err))
            else:
                print('amdgpu-plot has closed')
            env.GUT_CONST.PLOT = False

    # update gui
    for dk, dv in devices.items():
        for lk, lv in dv.items():
            if lk == 'card_num':
                data_value = 'card{}'.format(gpu_list.list[dk].get_params_value('card_num'))[:16]
            else:
                data_value = str(gpu_list.list[dk].get_params_value(lk))[:16]
                if data_value == '-1':
                    data_value = ''
            lv.set_text(data_value)
            lv.set_width_chars(17)

    while Gtk.events_pending():
        Gtk.main_iteration_do(True)
    # SEMAPHORE ############
    UD_SEM.release()
    ########################


def refresh(refreshtime, update_data, gpu_list, devices, cmd, gmonitor):
    """
    Method called for monitor refresh.
    :param refreshtime:  Amount of seconds to sleep after refresh.
    :type refreshtime: int
    :param update_data: Function that does actual data update.
    :type update_data: Callable
    :param gpu_list: A gpuList object with all gpuItems
    :type gpu_list: gpuList
    :param devices: A dictionary linking Gui items with data.
    :type devices: dict
    :param cmd: Subprocess return from running plot.
    :type cmd: subprocess.Popen
    :param gmonitor:
    :type gmonitor: Gtk
    :return:
    """
    while True:
        if gmonitor.quit:
            print('Quitting...')
            Gtk.main_quit()
            sys.exit(0)
        GLib.idle_add(update_data, gpu_list, devices, cmd)
        tst = 0.0
        sleep_interval = 0.2
        while tst < refreshtime:
            time.sleep(sleep_interval)
            tst += sleep_interval


def main():
    """
    Flow for amdgpu-monitor.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--about', help='README', action='store_true', default=False)
    parser.add_argument('--gui', help='Display GTK Version of Monitor', action='store_true', default=False)
    parser.add_argument('--log', help='Write all monitor data to logfile', action='store_true', default=False)
    parser.add_argument('--plot', help='Open and write to amdgpu-plot', action='store_true', default=False)
    parser.add_argument('--ltz', help='Use local time zone instead of UTC', action='store_true', default=False)
    parser.add_argument('--sleep', help='Number of seconds to sleep between updates', type=int, default=2)
    parser.add_argument('--no_fan', help='do not include fan setting options', action='store_true', default=False)
    parser.add_argument('-d', '--debug', help='Debug output', action='store_true', default=False)
    parser.add_argument('--pdebug', help='Plot debug output', action='store_true', default=False)
    args = parser.parse_args()

    # About me
    if args.about:
        print(__doc__)
        print('Author: ', __author__)
        print('Copyright: ', __copyright__)
        print('Credits: ', __credits__)
        print('License: ', __license__)
        print('Version: ', __version__)
        print('Maintainer: ', __maintainer__)
        print('Status: ', __status__)
        sys.exit(0)

    env.GUT_CONST.DEBUG = args.debug
    env.GUT_CONST.PDEBUG = args.pdebug
    if args.ltz:
        env.GUT_CONST.USELTZ = True
    if args.no_fan:
        env.GUT_CONST.show_fans = False
    if int(args.sleep) > 0:
        env.GUT_CONST.SLEEP = int(args.sleep)
    else:
        print('Invalid value for sleep specified.  Must be an integer great than zero')
        sys.exit(-1)

    if env.GUT_CONST.check_env() < 0:
        print('Error in environment. Exiting...')
        sys.exit(-1)

    # Get list of AMD GPUs and get basic non-driver details
    gpu_list = gpu.GpuList()
    gpu_list.set_gpu_list()

    # Check list of GPUs
    num_gpus = gpu_list.num_vendor_gpus()
    print('Detected GPUs: ', end='')
    for i, (k, v) in enumerate(num_gpus.items()):
        if i:
            print(', {}: {}'.format(k, v), end='')
        else:
            print('{}: {}'.format(k, v), end='')
    print('')
    if 'AMD' in num_gpus.keys():
        env.GUT_CONST.read_amd_driver_version()
        print('AMD: {}'.format(gpu_list.wattman_status()))
    if 'NV' in num_gpus.keys():
        print('nvidia smi: [{}]'.format(env.GUT_CONST.cmd_nvidia_smi))

    num_gpus = gpu_list.num_gpus()
    if num_gpus['total'] == 0:
        print('No GPUs detected, exiting...')
        sys.exit(-1)

    # Read data static/dynamic/info/state driver information for GPUs
    gpu_list.read_gpu_sensor_data(data_type='All')

    # Check number of readable/writable GPUs again
    num_gpus = gpu_list.num_gpus()
    print('{} total GPUs, {} rw, {} r-only, {} w-only\n'.format(num_gpus['total'], num_gpus['rw'],
                                                                num_gpus['r-only'], num_gpus['w-only']))

    time.sleep(1)
    # Generate a new list of only compatible GPUs
    com_gpu_list = gpu_list.list_gpus(compatibility='readable')

    if args.log:
        env.GUT_CONST.LOG = True
        env.GUT_CONST.log_file = './log_monitor_{}.txt'.format(
            env.GUT_CONST.now(ltz=env.GUT_CONST.USELTZ).strftime('%m%d_%H%M%S'))
        env.GUT_CONST.log_file_ptr = open(env.GUT_CONST.log_file, 'w', 1)
        gpu_list.print_log_header(env.GUT_CONST.log_file_ptr)

    if args.plot:
        args.gui = True
    if args.gui:
        # Display Gtk style Monitor
        devices = {}
        gmonitor = MonitorWindow(com_gpu_list, devices)
        gmonitor.connect('delete-event', gmonitor.set_quit)
        gmonitor.show_all()

        cmd = None
        if args.plot:
            env.GUT_CONST.PLOT = True
            if os.path.isfile('/usr/bin/amdgpu-plot'):
                plot_util = '/usr/bin/amdgpu-plot'
            else:
                plot_util = os.path.join(env.GUT_CONST.repository_path, 'amdgpu-plot')
            if os.path.isfile(plot_util):
                if env.GUT_CONST.PDEBUG:
                    cmd_str = '{} --debug --stdin --sleep {}'.format(plot_util, env.GUT_CONST.SLEEP)
                else:
                    cmd_str = '{} --stdin --sleep {}'.format(plot_util, env.GUT_CONST.SLEEP)
                cmd = subprocess.Popen(shlex.split(cmd_str), bufsize=-1, shell=False, stdin=subprocess.PIPE)
                com_gpu_list.print_plot_header(cmd.stdin)

        # Start thread to update Monitor
        _ = threading.Thread(target=refresh, daemon=True,
                             args=[env.GUT_CONST.SLEEP, update_data, com_gpu_list, devices, cmd, gmonitor]).start()

        Gtk.main()
    else:
        # Display text style Monitor
        try:
            while True:
                com_gpu_list.read_gpu_sensor_data(data_type='DynamicM')
                com_gpu_list.read_gpu_sensor_data(data_type='StateM')
                if not env.GUT_CONST.DEBUG: os.system('clear')
                if env.GUT_CONST.LOG:
                    print('{}Logging to:  {}{}'.format('\033[31m \033[01m', env.GUT_CONST.log_file, '\033[0m'))
                    com_gpu_list.print_log(env.GUT_CONST.log_file_ptr)
                com_gpu_list.print_table()
                time.sleep(env.GUT_CONST.SLEEP)
                if MonitorWindow.quit:
                    sys.exit(-1)
        except KeyboardInterrupt:
            if env.GUT_CONST.LOG:
                env.GUT_CONST.log_file_ptr.close()
            sys.exit(0)


if __name__ == '__main__':
    main()
