#!/usr/bin/env python3
""" amdgpu-chk  -  Checks OS/Python compatibility

    This utility verifies if the environment is compatible with amdgpu-utils.

    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-chk'
__version__ = 'v2.6.0'
__maintainer__ = 'RueiKe'
__status__ = 'Stable Release'

import argparse
import re
import subprocess
import os
import shlex
import platform
import sys
import shutil
import warnings
warnings.filterwarnings('ignore')


class GUT_CONST:
    def __init__(self):
        self.DEBUG = False

    @staticmethod
    def check_env():
        ret_val = [0, 0, 0]
        # Check python version
        required_pversion = [3, 6]
        (python_major, python_minor, python_patch) = platform.python_version_tuple()
        print('Using python ' + python_major + '.' + python_minor + '.' + python_patch)
        if int(python_major) < required_pversion[0]:
            print('          ' + '\x1b[1;37;41m' + ' but amdgpu-utils requires python ' +
                  str(required_pversion[0]) + '.' + str(required_pversion[1]) + ' or newer.' + '\x1b[0m')
            ret_val[0] = -1
        elif int(python_major) == required_pversion[0] and int(python_minor) < required_pversion[1]:
            print('          ' + '\x1b[1;37;41m' + ' but amdgpu-utils requires python ' +
                  str(required_pversion[0]) + '.' + str(required_pversion[1]) + ' or newer.' + '\x1b[0m')
            ret_val[0] = -1
        else:
            print('          ' + '\x1b[1;37;42m' + ' Python version OK. ' + '\x1b[0m')
            ret_val[0] = 0

        # Check Linux Kernel version
        required_kversion = [4, 8]
        linux_version = platform.release()
        print('Using Linux Kernel ' + str(linux_version))
        if int(linux_version.split('.')[0]) < required_kversion[0]:
            print('          ' + '\x1b[1;37;41m' + ' but amdgpu-util requires ' +
                  str(required_kversion[0]) + '.' + str(required_kversion[1]) + ' or newer.' + '\x1b[0m')
            ret_val[1] = -2
        elif int(linux_version.split('.')[0]) == required_kversion[0] and \
                int(linux_version.split('.')[1]) < required_kversion[1]:
            print('          ' + '\x1b[1;37;41m' + ' but amdgpu-util requires ' +
                  str(required_kversion[0]) + '.' + str(required_kversion[1]) + ' or newer.' + '\x1b[0m')
            ret_val[1] = -2
        else:
            print('          ' + '\x1b[1;37;42m' + ' OS kernel OK. ' + '\x1b[0m')
            ret_val[1] = 0

        # Check for amdgpu driver
        required_aversion = 'driver=amdgpu'
        lshw_out = subprocess.check_output(shlex.split('lshw -c video'), shell=False,
                                           stderr=subprocess.DEVNULL).decode().split('\n')
        config_found = False
        driver_str = ''
        for lshw_line in lshw_out:
            searchObj = re.search('configuration:', lshw_line)
            if searchObj:
                config_found = True
                lineitems = lshw_line.split(sep=':')
                driver_str = lineitems[1].strip()
                searchObj = re.search(required_aversion, driver_str)
                break
        if not config_found:
            driver_str = 'not detected'
            print('AMD GPU driver not detected')
            print('          ' + '\x1b[1;37;41m' + ' AMD\'s \'amdgpu\' driver package is required. ' + '\x1b[0m')
            ret_val[2] = -3
        if searchObj:
            print('AMD GPU driver is ' + driver_str)
            print('          ' + '\x1b[1;37;42m' + ' AMD driver OK. ' + '\x1b[0m')
            ret_val[2] = 0
        else:
            print('AMD GPU driver is ' + driver_str)
            print('          ' + '\x1b[1;37;41m' + ' AMD\'s \'amdgpu\' driver package is required. ' + '\x1b[0m')
            ret_val[2] = -3

        return ret_val


gut_const = GUT_CONST()


def is_venv_installed():
    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)
        output, error = p.communicate()
        # print('subprocess output: ', output.decode(), 'subprocess error: ', error.decode())
        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
    except:
        pass
    print('python3 venv is NOT installed')
    print('          ' + '\x1b[1;30;43m' + ' Python3 venv package \'python3-venv\' package is recommended. ' +
          '\x1b[0m')
    return False


def does_amdgpu_utils_env_exist():
    env_name = './amdgpu-utils-env/bin/activate'

    if os.path.isfile(env_name):
        print('amdgpu-utils-env available')
        print('          ' + '\x1b[1;37;42m' + ' amdgpu-utils-env OK. ' + '\x1b[0m')
        return True
    print('amdgpu-utils-env is NOT available')
    print('          ' + '\x1b[1;30;43m' + ' amdgpu-utils-env should be configured per User Guide. ' + '\x1b[0m')
    return False

def is_in_venv():
    python_path = shutil.which('python')

    if re.fullmatch(r'.*amdgpu-utils-env.*', python_path):
        print('In amdgpu-utils-env')
        print('          ' + '\x1b[1;37;42m' + ' amdgpu-utils-env is activated. ' + '\x1b[0m')
        return True
    print('Not in amdgpu-utils-env')
    print('          ' + '\x1b[1;30;43m' + ' amdgpu-utils-env should be activated per User Guide. ' + '\x1b[0m')
    return False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--about', help='README', action='store_true', default=False)
    parser.add_argument('-d', '--debug', help='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)

    gut_const.DEBUG = args.debug

    if gut_const.check_env() != [0, 0, 0]:
        print('Error in environment. Exiting...')
        sys.exit(-1)

    if not is_venv_installed() or not does_amdgpu_utils_env_exist():
        print('Environment not configured. WARNING')

    if not is_in_venv():
        print('Virtual Environment not activated. WARNING')


if __name__ == '__main__':
    main()
