#!python
""" gpu-chk  -  Checks OS/Python compatibility

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

    This utility verifies if the environment is compatible with *rickslab-gpu-utils*.

    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__ = 'RueiKe'
__copyright__ = 'Copyright (C) 2019 RicksLab'
__credits__ = ['Craig Echt - Testing, Debug, Verification, and Documentation']
__license__ = 'GNU General Public License'
__program_name__ = 'gpu-chk'
__maintainer__ = 'RueiKe'
__docformat__ = 'reStructuredText'
# pylint: disable=multiple-statements
# pylint: disable=line-too-long
# pylint: disable=bad-continuation

import argparse
import re
import subprocess
import os
import shlex
import platform
import sys
import shutil
import warnings
from GPUmodules import __version__, __status__

warnings.filterwarnings('ignore')
ENV_DIR = 'rickslab-gpu-utils-env'


class GutConst:
    """
    Base object for chk util.  These are simplified versions of what are in env module designed to run in python2
    in order to detect setup issues even if wrong version of python.
    """
    _verified_distros = ['Debian', 'Ubuntu', 'Gentoo', 'Arch']
    _dpkg_tool = {'Debian': 'dpkg', 'Ubuntu': 'dpkg', 'Arch': 'pacman', 'Gentoo': 'portage'}

    def __init__(self):
        self.DEBUG = False

    def check_env(self) -> dict:
        """
        Checks python version, kernel version, distro, and amd gpu driver version.

        :return:  A list of 4 integers representing status of 3 check items.
        """
        ret_val = {'python': 0, 'kernel': 0, 'system': 0, 'distribution': 0, 'driver': 0}
        # Check python version
        required_pversion = [3, 6]
        (python_major, python_minor, python_patch) = platform.python_version_tuple()
        print('Using python {}.{}.{}'.format(python_major, python_minor, python_patch))
        if int(python_major) < required_pversion[0]:
            print('          \x1b[1;37;41m but rickslab-gpu-utils requires python {}.{} or newer\x1b[0m'.format(
                  required_pversion[0], required_pversion[1]))
            ret_val['python'] = False
        elif int(python_major) == required_pversion[0] and int(python_minor) < required_pversion[1]:
            print('          \x1b[1;37;41m but rickslab-gpu-utils requires python {}.{} or newer \x1b[0m'.format(
                  required_pversion[0], required_pversion[1]))
            ret_val['python'] = False
        else:
            print('          \x1b[1;37;42m Python version OK. \x1b[0m')
            ret_val['python'] = True

        # Check Linux Kernel version
        required_kversion = [4, 8]
        linux_version = platform.release()
        print('Using Linux Kernel: {}'.format(linux_version))
        if int(linux_version.split('.')[0]) < required_kversion[0]:
            print('          \x1b[1;37;41m but rickslab-gpu-utils requires {}.{} or newer \x1b[0m'.format(
                  required_kversion[0], required_kversion[1]))
            ret_val['kernel'] = False
        elif int(linux_version.split('.')[0]) == required_kversion[0] and \
                int(linux_version.split('.')[1]) < required_kversion[1]:
            print('          \x1b[1;37;41m but rickslab-gpu-utils requires {}.{} or newer. \x1b[0m'.format(
                  required_kversion[0], required_kversion[1]))
            ret_val['kernel'] = False
        else:
            print('          \x1b[1;37;42m OS kernel OK. \x1b[0m')
            ret_val['kernel'] = True

        # Check Linux System Type
        init_type = 'Unknown'
        cmd_init = '/sbin/init' if os.path.isfile('/sbin/init') else shutil.which('init')
        if cmd_init:
            if os.path.islink(cmd_init):
                sys_path = os.readlink(cmd_init)
                init_type = 'systemd' if 'systemd' in sys_path else sys_path
            print('Using system type: {}'.format(init_type))
        if init_type == 'systemd':
            print('          \x1b[1;37;42m System type has been Validated. \x1b[0m')
            ret_val['system'] = True
        else:
            print('          \x1b[1;30;43m System type has not been verified. \x1b[0m')
            ret_val['system'] = False

        # Check Linux Distribution
        cmd_lsb_release = shutil.which('lsb_release')
        print('Using Linux distribution: ', end='')
        if cmd_lsb_release:
            distributor = description = None
            lsbr_out = subprocess.check_output(shlex.split('{} -a'.format(cmd_lsb_release)),
                                               shell=False, stderr=subprocess.DEVNULL).decode().split('\n')
            for lsbr_line in lsbr_out:
                if re.search('Distributor ID', lsbr_line):
                    lsbr_item = re.sub(r'Distributor ID:\s*', '', lsbr_line)
                    distributor = lsbr_item.strip()
                if re.search('Description', lsbr_line):
                    lsbr_item = re.sub(r'Description:\s*', '', lsbr_line)
                    if self.DEBUG: print('Distro Description: {}'.format(lsbr_item))
                    description = lsbr_item.strip()

            if distributor:
                print(description)
                if distributor in GutConst._verified_distros:
                    print('          \x1b[1;37;42m Distro has been Validated. \x1b[0m')
                    ret_val['distribution'] = True
                else:
                    print('          \x1b[1;30;43m Distro has not been verified. \x1b[0m')
                    ret_val['distribution'] = False
        else:
            print('unknown')
            print('          \x1b[1;30;43m [lsb_release] executable not found. \x1b[0m')
            ret_val['distribution'] = True

        # Check for amdgpu driver
        ret_val['driver'] = True if self.read_amd_driver_version() else True  # Ignore False
        return ret_val

    def read_amd_driver_version(self) -> bool:
        """
        Read the AMD driver version and store in GutConst object.

        :return: True if successful
        """
        cmd_dpkg = shutil.which('dpkg')
        if not cmd_dpkg:
            print('Command dpkg not found. Can not determine amdgpu version.')
            print('          \x1b[1;30;43m rickslab-gpu-utils can still be used. \x1b[0m')
            return True
        for pkgname in ['amdgpu', 'amdgpu-core', 'amdgpu-pro', 'rocm-utils']:
            try:
                dpkg_out = subprocess.check_output(shlex.split(cmd_dpkg + ' -l ' + pkgname),
                                                   shell=False, stderr=subprocess.DEVNULL).decode().split('\n')
            except (subprocess.CalledProcessError, OSError):
                continue
            for dpkg_line in dpkg_out:
                for driverpkg in ['amdgpu', 'rocm']:
                    if re.search(driverpkg, dpkg_line):
                        if self.DEBUG: print('Debug: ' + dpkg_line)
                        dpkg_items = dpkg_line.split()
                        if len(dpkg_items) > 2:
                            if re.fullmatch(r'.*none.*', dpkg_items[2]): continue
                            print('AMD: {} version: {}'.format(driverpkg, dpkg_items[2]))
                            print('          \x1b[1;37;42m AMD driver OK. \x1b[0m')
                            return True
        print('amdgpu/rocm version: UNKNOWN')
        print('          \x1b[1;30;43m rickslab-gpu-utils can still be used. \x1b[0m')
        return False


GUT_CONST = GutConst()


def is_venv_installed() -> bool:
    """
    Check if a venv is being used.

    :return: True if using venv
    """
    cmdstr = 'python3 -m venv -h > /dev/null'
    try:
        p = subprocess.Popen(shlex.split(cmdstr), shell=False, stdin=subprocess.PIPE,
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except (subprocess.CalledProcessError, OSError):
        pass
    else:
        output, _error = p.communicate()
        if not re.fullmatch(r'.*No module named.*', output.decode()):
            print('python3 venv is installed')
            print('          \x1b[1;37;42m python3-venv OK. \x1b[0m')
            return True
    print('python3 venv is NOT installed')
    print('          \x1b[1;30;43m Python3 venv package \'python3-venv\' is recommended for developers \x1b[0m')
    return False


def does_gpu_utils_env_exist() -> bool:
    """
    Check if venv exists.

    :return:  Return True if venv exists.
    """
    env_name = os.path.join('.', ENV_DIR, 'bin/activate')

    if os.path.isfile(env_name):
        print('{} available'.format(ENV_DIR))
        print('          \x1b[1;37;42m {} OK. \x1b[0m'.format(ENV_DIR))
        return True
    print('{} is NOT available'.format(ENV_DIR))
    print('          \x1b[1;30;43m {} can be configured per User Guide. \x1b[0m'.format(ENV_DIR))
    return False


def is_in_venv() -> bool:
    """
    Check if execution is from within a venv.

    :return: True if in venv
    """
    python_path = shutil.which('python3')
    if not python_path:
        print('Maybe python version compatibility issue.')
    else:
        if re.search(ENV_DIR, python_path):
            print('In {}'.format(ENV_DIR))
            print('          \x1b[1;37;42m {} is activated. \x1b[0m'.format(ENV_DIR))
            return True
        print('Not in {}, (Only needed if you want to duplicate development env)'.format(ENV_DIR))
        print('          \x1b[1;30;43m {} can be activated per User Guide. \x1b[0m'.format(ENV_DIR))
    return False


def main() -> None:
    """
    Main flow for chk utility.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--about', help='README', 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('Maintainer: ', __maintainer__)
        print('Status: ', __status__)
        sys.exit(0)

    system_status = GUT_CONST.check_env()
    if False in system_status.values():
        print(system_status)
        print('Error in environment. Exiting...')
        sys.exit(-1)

    if not is_venv_installed() or not does_gpu_utils_env_exist():
        print('Virtual Environment not configured. Only required by developers.')

    if not is_in_venv():
        pass
    print('')


if __name__ == '__main__':
    main()
