#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import logging
import argparse
import atexit
ANCERT_HOME_DIR = os.path.dirname(os.path.realpath(__file__))
sys.path.append(ANCERT_HOME_DIR)
from lib.utils import *
from lib.device import DeviceTree, Device
from lib.config import TESTS_DIR
from lib.expection import NonAvaliableController, NonAvaliableDevice

SUPPORT_COMPONENT = ['System', 'BIOS', 'CPU', 'Memory', 'Video', 'GPU', 'Storage',
                     'Network', 'NVMe', 'FC', 'IPMI', 'RAID', 'Kdump', 'Disk',
                     'Suspend', 'Misc']
RESULT, SAVE_LOG = True, False

def main(args):
    logging.info('Welcome to use the Ancert!')

    logging.info('Starting to probe devices.')
    dstree = DeviceTree()

    logging.info('Collecting host hardware information.')
    collect_host_info(logdir)

    logging.info('Loading testcase meta data.')
    cases = load_tests_yaml(TESTS_DIR)

    global SAVE_LOG
    if args.list_testcase:
        print('Testcase List:')
        name_max_length_lst = []
        for cur_cases in cases.items():
            for cur_case_info in cur_cases[1].items():
                for cur_info in cur_case_info[1]['cases']:
                    name_max_length_lst.append(len(cur_info['name']))
        name_max_length = max(name_max_length_lst) + 5
        print('{:<12}{:{}}{:<12}{:<12}{:<12}{:<12}{:<12}{:<12}'.format('Category:', 'Case:', name_max_length,
              'Mode:', 'Runtime:', 'Feature:', 'Required:', 'FileType:', 'FilePath:'))
        print('-'*(name_max_length + 84))
        for mode, val in cases.items():
            for name, cur_cases in val.items():
                for case in cur_cases['cases']:
                    print('{:<12}{:{}}{:<12}{:<12}{:<12}{:<12}{:<12}{:<12}'.format(cur_cases['category'],
                          case['name'], name_max_length, mode, case['runtime'], ','.join(case['feature']),
                          ','.join(case['required']), case['filetype'],
                          os.path.join(mode, name, case['filepath'])))
        return 0
    else:
        if not args.list_hardware:
            if args.category != 'System' and args.category.lower() not in cases[args.mode]:
                print('%s mode include:\n%s' % (args.mode, ', '.join([val['category']
                      for _, val in cases[args.mode].items()])))
                print('')
                print('%s mode does not have %s!' % (args.mode, args.category))
                return 1

    selected_components = []
    total_finished_tasks = []
    logging.debug('user selected %s test' % args.category)
    categories = [args.category]
    if args.category == 'System':
        categories = ['CPU', 'Memory', 'Storage', 'Network']
        if 'DISPLAY' in os.environ:
            categories.append('Video')
    if args.list_hardware:
        if args.list_hardware == 'All':
            logging.debug('user selected list %s hardware' % args.list_hardware)
            categories = SUPPORT_COMPONENT
            categories.remove('System')
        elif args.list_hardware == 'System':
            categories = ['CPU', 'Memory', 'Storage', 'Network']
            if 'DISPLAY' in os.environ:
                categories.append('Video')
        else:
            categories = [args.list_hardware]

    if args.cases:
        for category in categories:
            for case in args.cases:
                if case not in [c['name'] for c in cases[args.mode][category.lower()]['cases']]:
                    print('Does not have test case %s for %s' % (case, ', '.join(categories)))
                    return 1
        for category in categories:
            for case in [c for c in cases[args.mode][category.lower()]['cases']]:
                if case['name'] not in args.cases:
                    cases[args.mode][category.lower()]['cases'].remove(case)

    if args.mode in ['function', 'stress']:
        cases = cases[args.mode]
        if args.mode == 'stress':
            stress_category = ['Network', 'NVMe', 'CPU', 'Memory']
            if args.category:
                if args.category not in stress_category:
                    print('Error: Please select category correctly from {}'.format(stress_category))
                    sys.exit(1)
            else:
                categories = stress_category
    else:
        cases = cases['certify']

    if args.feature:
        for category in categories:
            for _ in range(len(cases[args.mode][category.lower()]['cases'])):
                case = cases[args.mode][category.lower()]['cases'].pop(0)
                if args.feature in case['feature']:
                    cases[args.mode][category.lower()]['cases'].append(case)

    print('\n\t\t\tWelcome to use the Ancert!\n')

    if args.list_hardware:
        print('OS information:')
        print('-'*9)
    get_os_version()
    get_kernel_version()
    print('')
    SAVE_LOG = True

    for comp_str in categories:
        try:
            loc = locals()
            exec('from lib.component import %s as component' % comp_str)
            component = loc['component']
            comp_inst = component(logdir, cases.get(comp_str.lower(), {}), comp_str, args)
            if args.category == 'System' or args.list_hardware == 'System':
                comp_inst.is_system_test = True
            logging.debug('starting to init component %s' % comp_str)
            selected_components.append(comp_inst)
            comp_inst.build(dstree)
        except (NonAvaliableController, NonAvaliableDevice) as e:
            if not args.list_hardware:
                print(e)
                return 1
        except NonDriverAttached as e:
            if args.category == 'System':
                print(e)
                return 1
    else:
        save_device_info(selected_components, dstree, args, logdir)

    if args.list_hardware:
        print('Hardware Information:')
        print('-'*9)
        if args.list_hardware == 'System':
            print('System:\n  %s' % get_system_info())
        for comp_inst in selected_components:
            if comp_inst.category in ['Kdump', 'Disk', 'Suspend', 'Misc']:
                continue
            if comp_inst.category in ['Memory']:
                print('%s:' % comp_inst.category)
                for phy in comp_inst.physical_memory:
                    print('  {:<4}{:50}'.format('[-]', '%s [%s]' % (phy[0], phy[1])))
                continue
            if comp_inst.category in ['IPMI'] and not comp_inst.fake_device_name:
                continue
            if comp_inst.available or comp_inst.unavailable:
                print('%s:' % comp_inst.category)
            if comp_inst.category in ['BIOS']:
                print('  {:<4}{:50}'.format('[-]', '%s' % comp_inst.available[0].name.strip()))
                continue
            for index, ctr in enumerate(comp_inst.available):
                if comp_inst.category in ['CPU', 'Video']:
                    print('  {:<4}{:50}'.format('[-]', '%s [%s]' % (ctr.name.strip(), ctr.driver if ctr
                                                .driver else NO_DRIVER_STR)))
                    continue
                if comp_inst.category in ['IPMI']:
                    print('  {:<4}{:50}'.format('[-]', '%s [%s]' % (comp_inst.fake_device_name,
                              comp_inst.available[0].available_devices[0].driver)))
                    continue
                print('  {:<4}{:50}'.format('[%s]' % (index+1), '%s [%s]' % (ctr.name, ctr
                                            .driver if ctr.driver else NO_DRIVER_STR)))
                if comp_inst.category in ['Storage', 'NVMe', 'FC', 'RAID']:
                    for device in ctr.devices:
                        if device.is_partition:
                            continue
                        print('        {:<4}{:50}'.format('[-]' if device in ctr
                              .available_devices else '[-]', '%s [%s]'
                              % (device.get_property('DRIVE_FULL_NAME'), device.get_property('DISK_SIZE'))))
            for ctr in comp_inst.unavailable:
                print('  {:<4}{:50}'.format('[*]', '%s [%s]' % (ctr.name, ctr.driver if ctr
                                            .driver else NO_DRIVER_STR)))
                if comp_inst.category in ['Storage', 'NVMe', 'FC', 'RAID']:
                    for device in ctr.devices:
                        print('        {:<4}{:50}'.format('[-]', '%s [%s]'
                              % (device.get_property('DRIVE_FULL_NAME'), device.get_property('DISK_SIZE'))))
        return 0

    c = 0
    is_fake_dev = False
    for comp_inst in selected_components:
        printed = []
        ins_name_lst = [str(ins).split("'")[1].split('.')[-1] for ins in list(comp_inst.__class__.__bases__)]
        for gname, tasks in comp_inst.next_tasks():
            if tasks[0].test.controller not in printed:
                if c == 0 and 'FakeComponent' not in ins_name_lst:
                    print('')
                if 'fake device' not in tasks[0].test.controller.name:
                    if comp_inst.name in ['memory']:
                        print('Tested controller: %s' % tasks[0].test.controller.name)
                    else:
                        print('Tested controller: %s [%s]' % (tasks[0].test.controller.name.strip(),
                                                                tasks[0].test.controller.driver))
                else:
                    is_fake_dev = True
                printed.append(tasks[0].test.controller)
            run_task(gname, tasks)
            c += 1
            total_finished_tasks.extend(tasks)
    if not is_fake_dev:
        print('')

    global RESULT
    if total_finished_tasks:
        logging.info('all tests finished')
        print('Results Summary:')
        test_name_max_len = max([len(t.name) for t in total_finished_tasks]) + 5
        print('{:<12}{:{}}{:<12}{:<52}'.format('Category:', 'Case:', test_name_max_len, 'Result:', 'Device name:'))
        print('-'*(test_name_max_len + 36))
        for t in total_finished_tasks:
            if t.result() == 0:
                ret = 'PASS'
            elif t.result() == 2:
                ret = 'SKIP'
            else:
                ret = 'FAIL'
                RESULT = False
            logging.info('test case %s %sed on %s' % (t.name, ret, t.test.controller.name.strip()))
            print('{:<12}{:{}}{:<12}{:<52}'.format(t.test.controller.component.category,
                  t.name, test_name_max_len, ret, '' if 'fake device' in t.test.controller.name.strip()
                  else t.test.controller.name.strip()))
        print('-'*(test_name_max_len + 36))
    else:
        print('No testcase selected, please double check the options!')
        return 1

    if args.report:
        generate_summary_report(args, total_finished_tasks, logdir)

    if RESULT and save_test_result(selected_components, args, logdir, RESULT):
        return 0
    return 1


if '__main__' == __name__:
    os.chdir(ANCERT_HOME_DIR)
    if os.getuid() > 0:
        print('You need to be root to run this program.\n')
        sys.exit(1)

    args_parser = argparse.ArgumentParser(description='ancert',
                                          formatter_class=argparse.RawDescriptionHelpFormatter,
                                          epilog='Example:\n  python3 ancert -h\n  '
                                                 'python3 ancert -g Storage -c fs_rw raw_disk_rw')
    args_parser.add_argument('-g', '--category', type=str, default='', choices=SUPPORT_COMPONENT,
                             help='Category')
    args_parser.add_argument('-f', '--feature', type=str, default='', help='Feature')
    args_parser.add_argument('-c', '--cases', type=str, help='Cases', nargs='+')
    args_parser.add_argument('-i', '--index', type=int, default=1, help='Want to test device index')
    args_parser.add_argument('--lts_ip', type=str, help='LTS ip address')
    args_parser.add_argument('--exclude_interface', type=str, help='exclude interface name such as eth0')
    args_parser.add_argument('--single_mode', action='store_true', help='Run System test without LTS')
    args_parser.add_argument('-r', '--report', action='store_true', help='Summary report')
    args_parser.add_argument('--tone', action='store_true', help='Work with Tone')
    args_parser.add_argument('--list_testcase', action='store_true', help='List all the testcases')
    args_parser.add_argument('--list_hardware', type=str, default='', choices=['All', *SUPPORT_COMPONENT],
                             help='List hardware on the system for target category')
    args_parser.add_argument('-m', '--mode', type=str, default='certify', choices=['certify', 'function', 'stress'],
                             help='Currently supported test mode: certify, function, stress')
    args, unknown = args_parser.parse_known_args()
    if unknown:
        print('Unknown option: %s' % ', '.join(unknown))
        sys.exit(1)
    if args.list_hardware:
        if args.category:
            print('--category and --list_hardware can\'t use together!')
            sys.exit(1)
        if args.cases:
            print('--cases and --list_hardware can\'t use together!')
            sys.exit(1)
    else:
        if not args.list_testcase:
            if not args.category:
                print('Please select a category!')
                sys.exit(1)

    if args.single_mode and args.lts_ip:
        print('--lts_ip and --single_mode can\'t use together!')
        sys.exit(1)

    if args.mode not in ['certify'] and args.category in ['System']:
        print('System test cannot work with --mode certify!')
        sys.exit(1)

    if args.category in ['System']:
        if args.cases:
            print('--category and --cases can\'t use together!')
            sys.exit(1)

    if args.category in ['System', 'Suspend', 'Kdump']:
        if not args.single_mode and not args.lts_ip:
            print('Please input LTS ip address or run single mode for %s test' % args.category)
            sys.exit(1)

    if args.category in ['Network'] and not args.lts_ip:
       print('Please input LTS ip address for Network test')
       sys.exit(1)

    if args.category in ['System', 'Network']:
        if args.lts_ip:
            if not check_ipaddr_connection(args.lts_ip):
                sys.exit(1)
        if args.exclude_interface:
            exclude_interface_list = (args.exclude_interface).split(',')
            exclude_file_path = ANCERT_HOME_DIR + '/etc/exclude_interface'
            file_handle = open(exclude_file_path, "w+")
            for single_exclude_interface in exclude_interface_list:
                if not check_interface(single_exclude_interface):
                    sys.exit(1)
                line = file_handle.writelines(single_exclude_interface + '\n')
            file_handle.close()

    logdir = logging_config(args)

    @atexit.register
    def trigger_before_exit():
        if SAVE_LOG:
            print('')
            pack_log_dir(logdir, args, RESULT)

    if args.single_mode:
        if args.category in ['System'] or args.list_hardware in ['System']:
            print('\n*NOTE*: You are using ancert single testing mode! Single testing mode can\'t fully verify your\n'
                  '        system compatibility with Anolis OS. If you want to run all tests, please configure LTS\n'
                  '        system, then add option --lts_ip [lts ip address] to ancert.')

    ret = main(args)
    if ret != 0:
        RESULT = False
    sys.exit(ret)