#!/usr/bin/env python3

################################################################################
# Module import section
################################################################################
import argparse
import subprocess
import sys


################################################################################
# Global variables and defaults
################################################################################
iointerfaces_allowed = [
    'async',
    'cached',
    'direct',
    'sync',
    'dsync'
  ]

ioping_arguments = {
    'location': '',
    'cmd': 'ioping --batch',
    'count': '10',
    'interval': '1s',
    'options': ''
  }

################################################################################
# Global functions
################################################################################
# Function to run the shell command and capture outputs and returncode
def capture_ioping_results():

   if ioping_arguments['location'] == '':
       return 1, '', 'No ioping_locationectory given.'

   cmd = '%s %s --count %s --interval %s %s' % (ioping_arguments['cmd'],
                                                ioping_arguments['options'],
                                                ioping_arguments['count'],
                                                ioping_arguments['interval'],
                                                ioping_arguments['location'])
   ioping_arguments['cmd'] = cmd
   ioping_output = subprocess.run(cmd.split(),
                                  stdout=subprocess.PIPE,
                                  stderr=subprocess.PIPE,
                                  encoding='utf-8')
   return ioping_output.returncode, ioping_output.stdout, ioping_output.stderr


# Function to output the ioping capture on an interactive shell
def ioping_output_shell():
    returncode, stdout, stderr =  capture_ioping_results()
    print("%s\n" % (ioping_arguments['cmd']))

    if returncode == 0:
        print(stdout)
    else:
        print(stderr)

    return returncode


# Function to ouptut the ioping capure for telegraf
def ioping_output_telegraf():
    returncode, stdout, stderr =  capture_ioping_results()

    if not returncode == 0:
        print(stderr)
        return returncode

    ioping_statistic_keys = [
            'statistics_count',
            'runtime',
            'iops',
            'transfer_speed',
            'min_request_time',
            'avg_request_time',
            'max_request_time',
            'std_deviation_request_time',
            'total_requests',
            'total_running_time'
            ]

    try:
        ioping_statistics = dict(zip(ioping_statistic_keys, stdout.split()))
    except Exception as e:
        print(e)
        returncode = 1

    ioping_statistics_string = ",".join(f'{key}={value}' for key,value in ioping_statistics.items())
    print("ioping,location=%s %s" %(ioping_arguments['location'], ioping_statistics_string))

    return returncode

def get_parser():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-c', '--count',
                        help = 'stop after <count> requests',
                        default = '10')
    parser.add_argument('-i', '--interval',
                        help = 'interval between requests',
                        default = '1s')
    parser.add_argument('--iointerface',
                        help = 'Which ionterface to use: async|cached|direct|sync|dsync',
                        default = '')
    parser.add_argument('-l', '--location',
                        help = 'directory|file|device to test. Best would be a directory to not shred anything accidentially',
                        default = '/tmp')
    parser.add_argument('-m', '--mode',
                        help = 'Which output to use: shell|telegraf',
                        default = 'telegraf')
    parser.add_argument('-w', '--read-write',
                        help = 'Use read-write test instead of read.',
                        action = 'store_true',
                        default = 'False')
    return parser


# Main function for shell use
def main():
    parser = get_parser()
    args = parser.parse_args()

    ioping_arguments['location'] = args.location
    ioping_arguments['count'] = args.count
    ioping_arguments['interval'] = args.interval
    if args.read_write == True:
        ioping_arguments['options'] += '--read-write '
    if not args.iointerface == '':
        if args.iointerface in iointerfaces_allowed:
            ioping_arguments['options'] += '--%s ' % args.iointerface


    if args.mode == 'shell':
        returncode = ioping_output_shell()
    elif args.mode == 'telegraf':
        returncode = ioping_output_telegraf()
    else:
        returncode = 1

    return returncode


################################################################################
# Main main main
################################################################################
if __name__ == '__main__':
    sys.exit(main())