#!/usr/bin/env python
# coding: utf-8

import argparse
import os
import sys


def get_variable_from_env(variable):
    env_result = os.environ.get(variable)
    if not env_result:
        raise EnvironmentError('Environment variable {} was not found'.format(variable))
    return env_result


def get_shard_info(var, display=True):
    result = ''
    env_shard = get_variable_from_env('LABELS_shard').split('_')
    if var == 'shard_min':
        result = env_shard[-2]
    elif var == 'shard_max':
        result = env_shard[-1]
    if display:
        print(result)
    return result


def get_cpu_info(var):
    cpu_var = get_variable_from_env(var.upper()).split('.')[0]
    print(cpu_var)
    return cpu_var



def get_mem_info(var):
    env_mem = get_variable_from_env(var.upper())
    print(env_mem)
    return env_mem


def parse_args(args):
    opts = argparse.ArgumentParser(prog='select_env_resource',
                                   formatter_class=argparse.RawTextHelpFormatter,
                                   description="""
DESCRIPTION
Tool for selecting important environment variables inside YP pods:
shard_min
shard_max
cpu_guarantee
cpu_limit
mem_guarantee
mem_limit

Feel free to add new variables in the future.
Example:
select_env_resource.py cpu_limit
""")

    opts.add_argument('-e', '--env', dest='variable_name', type=str, help='Specifies variable name',
                      choices=['cpu_guarantee', 'cpu_limit', 'shard_min', 'shard_max', 'mem_guarantee', 'mem_limit'])

    if args:
        opts = opts.parse_args(args)
    else:
        opts = opts.parse_args()
    return opts


if __name__ == '__main__':
    options = parse_args(sys.argv[1:])
    varname = options.variable_name
    if varname.startswith('cpu_'):
        get_cpu_info(varname)
    elif varname.startswith('mem_'):
        get_mem_info(varname)
    elif varname.startswith('shard_'):
        get_shard_info(varname)
