#!/usr/bin/env python3
""" gpu-mon  -  Displays a continuously updating view of the status of all
                active GPUs.

    Part of the rickslab-gpu-utils package which includes gpu-ls, gpu-mon,
    gpu-pac, and gpu-plot.

    A utility to give the current state of all compatible 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.  The *--verbose*
    option will display progress and informational messages generated by the
    utilities.

    Copyright (C) 2019  RicksLab

    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__ = 'RicksLab'
__copyright__ = 'Copyright (C) 2019 RicksLab'
__license__ = 'GNU General Public License'
__program_name__ = 'gpu-mon'
__maintainer__ = 'RicksLab'
__docformat__ = 'reStructuredText'
# pylint: disable=multiple-statements
# pylint: disable=line-too-long
# pylint: disable=consider-using-f-string

import argparse
import subprocess
import threading
import os
import logging
import sys
from shlex import split as shlex_split
import shutil
from time import sleep
import signal
from typing import Callable, Any


try:
    import gi
    gi.require_version('Gtk', '3.0')
    from gi.repository import GLib, Gtk
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)

from GPUmodules import __version__, __status__, __credits__
from GPUmodules import GPUgui
from GPUmodules import GPUmodule as Gpu
from GPUmodules import env

set_gtk_prop = GPUgui.GuiProps.set_gtk_prop
LOGGER = logging.getLogger('gpu-utils')


def ctrl_c_handler(target_signal: Any, _frame: Any) -> None:
    """
    Signal catcher for ctrl-c to exit monitor loop.

    :param target_signal: Target signal name
    :param _frame: Ignored
    """
    LOGGER.debug('ctrl_c_handler (ID: %s) has been caught. Setting quit flag...', target_signal)
    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: bool = False
    item_width: int = env.GUT_CONST.mon_field_width
    label_width: int = 12

    def __init__(self, gpu_list, devices):

        init_chk_value = Gtk.init_check(sys.argv)
        LOGGER.debug('init_check: %s', init_chk_value)
        if not init_chk_value[0]:
            print('Gtk Error, Exiting')
            sys.exit(-1)
        Gtk.Window.__init__(self, title=env.GUT_CONST.gui_window_title)
        self.set_border_width(0)
        self.set_resizable(False)
        GPUgui.GuiProps.set_style()

        if env.GUT_CONST.icon_file:
            LOGGER.debug('Icon file: [%s]', env.GUT_CONST.icon_file)
            if os.path.isfile(env.GUT_CONST.icon_file):
                self.set_icon_from_file(env.GUT_CONST.icon_file)

        grid = Gtk.Grid()
        self.add(grid)

        col = 0
        row = 0
        num_amd_gpus = gpu_list.num_gpus()['total']
        if env.GUT_CONST.debug:
            debug_label = Gtk.Label(name='warn_label')
            debug_label.set_markup('<big><b> DEBUG Logger Active </b></big>')
            lbox = Gtk.Box(spacing=6, name='warn_box')
            set_gtk_prop(debug_label, top=1, bottom=1, right=1, left=1)
            lbox.pack_start(debug_label, True, True, 0)
            grid.attach(lbox, 0, row, num_amd_gpus+1, 1)
        row += 1
        if env.GUT_CONST.log:
            log_label = Gtk.Label(name='warn_label')
            log_label.set_markup('<big><b> Logging to:    </b>{}</big>'.format(env.GUT_CONST.log_file))
            lbox = Gtk.Box(spacing=6, name='warn_box')
            set_gtk_prop(log_label, top=1, bottom=1, right=1, 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(name='white_label', halign=Gtk.Align.CENTER, valign=Gtk.Align.CENTER)}
        row_labels['card_num'].set_markup('<b>Card #</b>')
        for param_name, param_label in Gpu.GpuItem.table_param_labels.items():
            row_labels[param_name] = Gtk.Label(name='white_label', halign=Gtk.Align.START, valign=Gtk.Align.CENTER)
            row_labels[param_name].set_markup('<b>{}</b>'.format(param_label))
        for row_label_item in row_labels.values():
            lbox = Gtk.Box(spacing=6, name='head_box')
            set_gtk_prop(lbox, top=1, bottom=1, right=1, left=1)
            set_gtk_prop(row_label_item, top=1, bottom=1, right=4, left=4)
            lbox.pack_start(row_label_item, True, True, 0)
            grid.attach(lbox, col, row, 1, 1)
            row += 1

        for gpu in gpu_list.gpus():
            devices[gpu.prm.uuid] = {'card_num':  Gtk.Label(name='white_label')}
            devices[gpu.prm.uuid]['card_num'].set_markup('<b>CARD{}</b>'.format(gpu.get_params_value('card_num')))
            devices[gpu.prm.uuid]['card_num'].set_use_markup(True)
            for param_name in Gpu.GpuItem.table_param_labels:
                devices[gpu.prm.uuid][param_name] = Gtk.Label(label=gpu.get_params_value(str(param_name)),
                                                              name='white_label')
                devices[gpu.prm.uuid][param_name].set_width_chars(self.item_width)
                set_gtk_prop(devices[gpu.prm.uuid][param_name], width_chars=self.item_width)

        for gui_component in devices.values():
            col += 1
            row = row_start
            for comp_name, comp_item in gui_component.items():
                comp_item.set_text('')
                if comp_name == 'card_num':
                    lbox = Gtk.Box(spacing=6, name='head_box')
                else:
                    lbox = Gtk.Box(spacing=6, name='med_box')
                set_gtk_prop(lbox, top=1, bottom=1, right=1, left=1)
                set_gtk_prop(comp_item, top=1, bottom=1, right=3, left=3, width_chars=self.item_width)
                lbox.pack_start(comp_item, True, True, 0)
                grid.attach(lbox, col, row, 1, 1)
                row += 1

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


def update_data(gpu_list: Gpu.GpuList, devices: dict, cmd: subprocess.Popen) -> None:
    """
    Update monitor data with data read from GPUs.

    :param gpu_list: A gpuList object with all gpuItems
    :param devices: A dictionary linking Gui items with data.
    :param cmd: Subprocess return from running plot.
    """
    # SEMAPHORE ############
    if not UD_SEM.acquire(blocking=False):
        if env.GUT_CONST.verbose: print('Update while updating, skipping new update')
        LOGGER.debug('Update while updating, skipping new update')
        return
    ########################
    gpu_list.read_gpu_sensor_set(data_type=Gpu.GpuItem.SensorSet.Monitor)
    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:
            LOGGER.debug('gpu-plot has closed: [%s]', except_err)
            print('gpu-plot has closed')
            env.GUT_CONST.plot = False

    # update gui
    for uuid, gui_component in devices.items():
        for comp_name, comp_item in gui_component.items():
            if comp_name == 'card_num':
                comp_item.set_markup('<b>Card{}</b>'.format(gpu_list[uuid].get_params_value('card_num')))
            else:
                data_value_raw = gpu_list[uuid].get_params_value(comp_name)
                LOGGER.debug('raw data value: %s', data_value_raw)
                data_value_raw = Gpu.format_table_value(data_value_raw, comp_name)
                data_value = str(data_value_raw)[:MonitorWindow.item_width]
                comp_item.set_text(data_value)
            set_gtk_prop(comp_item, width_chars=MonitorWindow.item_width)

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


def refresh(refreshtime: int, update_data_func: Callable, gpu_list: Gpu.GpuList, devices: dict,
            cmd: subprocess.Popen, gmonitor: Gtk.Window) -> None:
    """
    Method called for monitor refresh.

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


def main() -> None:
    """
    Flow for gpu-mon.
    """
    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 gpu-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('--verbose', help='Display informational message of GPU util progress',
                        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: ', *['\n      {}'.format(item) for item in __credits__])
        print('License: ', __license__)
        print('Version: ', __version__)
        print('Install Type: ', env.GUT_CONST.install_type)
        print('Maintainer: ', __maintainer__)
        print('Status: ', __status__)
        sys.exit(0)

    if int(args.sleep) <= 1:
        print('Invalid value for sleep specified.  Must be an integer great than zero')
        sys.exit(-1)
    env.GUT_CONST.set_args(args, __program_name__)
    LOGGER.debug('########## %s %s', __program_name__, __version__)

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

    # Get list of GPUs and exit if no GPUs detected
    gpu_list = Gpu.GpuList()
    gpu_list.set_gpu_list()
    num_gpus = gpu_list.num_gpus()
    if num_gpus['total'] == 0:
        print('No GPUs detected, exiting...')
        sys.exit(-1)

    # Display vendor and driver details
    Gpu.print_driver_vendor_summary(gpu_list)

    # Read data static/dynamic/info/state driver information for GPUs
    gpu_list.read_gpu_sensor_set(data_type=Gpu.GpuItem.SensorSet.All)

    # Check number of readable/writable
    print('All GPUs:\n    {}'.format(gpu_list))

    # Select GPU's appropriate for monitor
    com_gpu_list = Gpu.set_mon_plot_compatible_gpu_list(gpu_list)

    # Check readable and compatible GPUs
    num_gpus = com_gpu_list.num_gpus()
    print('Compatible GPUs:')
    if num_gpus['total'] == 0:
        print('No readable and compatible GPUs detected, exiting...')
        sys.exit(-1)
    print('    {}'.format(com_gpu_list))

    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', buffering=1, encoding='utf-8')
        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 env.GUT_CONST.install_type == 'repository':
                plot_util = './gpu-plot'
            else:
                plot_util = shutil.which('gpu-plot')
            if not plot_util:
                plot_util = os.path.join(env.GUT_CONST.repository_path, 'gpu-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)
                # Do not use with, as cmd is meant to stay open as long as monitor is running.
                cmd = subprocess.Popen(shlex_split(cmd_str), bufsize=-1, shell=False, stdin=subprocess.PIPE)
                com_gpu_list.print_plot_header(cmd.stdin)
            else:
                print('Fatal Error: gpu-plot not found.')

        # 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_set(data_type=Gpu.GpuItem.SensorSet.Monitor)
                os.system('clear')
                if env.GUT_CONST.debug:
                    print('{}DEBUG logger is active{}'.format((env.GUT_CONST.mark_up_codes['red'] +
                                                               env.GUT_CONST.mark_up_codes['bold']),
                                                              env.GUT_CONST.mark_up_codes['reset']))
                if env.GUT_CONST.log:
                    print('{}Logging to: {}{}'.format((env.GUT_CONST.mark_up_codes['red'] +
                                                       env.GUT_CONST.mark_up_codes['bold']),
                                                      env.GUT_CONST.log_file,
                                                      env.GUT_CONST.mark_up_codes['reset']))
                    com_gpu_list.print_log(env.GUT_CONST.log_file_ptr)
                com_gpu_list.print_table()
                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()
