#!/usr/bin/python3

'''Driver package query/installation tool for UOS'''

#from lib2to3.pgen2 import driver
import os
from re import T
import click
import subprocess
from nvidia_detect import NVDetect


CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])


def _is_run_in_live():
    result = False
    with open('/proc/cmdline', mode='r', encoding='utf-8') as f:
        cmdline = f.read()
        for item in cmdline.split(" "):
            if item == 'boot=live':
                result = True
                break
        f.close()
    return result


class Config(object):
    def __init__(self):
        self.installer_init = False
        self.installer_target = False

pass_config = click.make_pass_decorator(Config, ensure=True)

@click.group(context_settings=CONTEXT_SETTINGS)
@click.option('--installer-init', is_flag=True, help='Used only during deepin-installer prepare phase.')
@click.option('--installer-target', is_flag=True, help='Used only in the deepin-installer phase.')
@pass_config
def greet(config, installer_init, installer_target, **kwargs):
    if not detect.get_nvidia_gpu_sysfs():
        print("Error: No NVIDIA GPU found.")
        exit(-1)
    if installer_init:
        if not _is_run_in_live():
            print("Error: the option '--installer-init' used only during deepin-installer prepare phase.")
            exit(-1)
        config.installer_init = True
    
    if installer_target:
        if not _is_run_in_live():
            print("Error: the option '--installer-target' used only in the deepin-installer phase.")
            exit(-1)
        config.installer_target = True


def auto_install(mini_install=False):
    drivers = detect.system_nvidia_driver_packages()
    recommnd_version = max(drivers.keys())
    detect.install_nvidia_driver(recommnd_version, mini_install)

def command_autoinstall(args):
    if args.installer_init:
        print("Nothing to do!")
    elif args.installer_target:
        auto_install(True)
    else:
        auto_install(False)


@greet.command()
@click.option('--installer-init', is_flag=True, help='Used only during deepin-installer prepare phase')
@click.option('--installer-target', is_flag=True, help='Used only in the deepin-installer phase')
@pass_config
def autoinstall(config, **kwargs):
    '''Install drivers that are appropriate for automatic installation.'''
    if kwargs.get('installer_init'):
        config.installer_init = True
    elif kwargs.get('installer_target'):
        config.installer_target = True
    command_autoinstall(config)


def command_getdev(args):
    sysfs_id = detect.get_nvidia_gpu_sysfs()
    module_alias = detect.get_nvidia_gpu_modalias()
    print("== %s ===" % sysfs_id)
    print('%-9s: %s' % ("modalias", module_alias))
    print('%-9s: %s\n' % ("vendor", "NVIDIA Corporation"))
    drivers = detect.system_nvidia_driver_packages()
    recommnd_version = max(drivers.keys())
    current_driver = detect.get_current_driver()
    for ver,drv in  drivers.items():
        str = '%-9s: %-18s %-9s %s' % ("driver", drv,
                                        "current" if drv == current_driver else "-------",
                                        "recommend" if ver == recommnd_version else "---------")
        print(str)

@greet.command()
@pass_config
def getdev(config, **kwargs):
    '''Get nvidia GPU device info by specific driver version. Such as 390 (Corresponds to nvidia-driver-390).'''
    command_getdev(config)

def command_getjson(args):
    json_data = detect.get_support_json(args.version)
    print(json_data)

@greet.command()
@click.argument('version')
@pass_config
def getjson(config, **kwargs):
    '''Get support GPUs info　by specific driver version. Such as 390 (Corresponds to nvidia-driver-390).'''

    if kwargs.get('version'):
        config.version = ''.join(kwargs.get('version'))
        command_getjson(config)


def command_install(args):
    detect.install_nvidia_driver(args.version, mini_install=False)

@greet.command()
@click.argument('version')
@pass_config
def install(config, **kwargs):
    '''Install the nvidia driver by specified version. Such as 390 (Corresponds to nvidia-driver-390).'''

    if kwargs.get('version'):
        config.version = ''.join(kwargs.get('version'))
        command_install(config)

if __name__ == '__main__':
    detect = NVDetect()
    greet()