#! /usr/bin/python3
#-------------------------------------------------------------------------------
#
# Copyright (C) 2020-2021 NVIDIA Corporation. ALL RIGHTS RESERVED.
# Copyright 2014,2015,2016,2017,2018, Cumulus Networks Inc  all rights reserved
#
#-------------------------------------------------------------------------------
#

#-------------------------------------------------------------------------------
#
# Imports
#

import warnings
import re
import os
import sys
import signal
import traceback
import argparse
import pickle
import json

# Classes
#
class ArgParseError(RuntimeError):
    pass

class SMONDRuntimeError(RuntimeError):
    pass

#-------------------------------------------------------------------------------
#
# Functions
#

#--------------------
#
# warning formats
#

def smondwarn(message, category, filename, lineno, line=None):
    return '%s:%s : %s : %s\n' % (filename, lineno, category.__name__, message)

#--------------------
#
# normal exit
#
def exit_normally(signum=0, frame=None):
    syslog.syslog(syslog.LOG_INFO, "exiting normally")
    sys.stderr.write("%s : exiting normally\n" % sys.argv[0])
    exit(0)

path = '/run/cache/cumulus/unit_state'
def unpickle(f):
    try:
        p_path = os.path.join(path,f, 'dump')
        f = open(p_path, "rb")
        m = pickle.load(f)
        f.close()
    except:
        return None
    return m

def print_temp(m):

    t_in = None; t_max = None; t_min = None; t_crit = None
    t_lcrit = None; t_max_hyst = None; t_fan_min = None
    t_fan_max = None

    for key,value in m.items():
        if value == None:
            continue
        if re.match('temp[0-9]*_input', key):
            t_in = float(value)
        elif re.match('temp[0-9]*_fan_max', key):
            t_fan_max = value
        elif re.match('temp[0-9]*_fan_min', key):
            t_fan_min = value
        elif re.match('temp[0-9]*_max_alarm', key):
            continue
        elif re.match('temp[0-9]*_min_alarm', key):
            continue
        elif re.match('temp[0-9]*_crit_alarm', key):
            continue
        elif re.match('temp[0-9]*_max', key):
            t_max = value
        elif re.match('temp[0-9]*_min', key):
            t_min = value
        elif re.match('temp[0-9]*_crit_hyst', key):
            continue
        elif re.match('temp[0-9]*_crit', key):
            t_crit = value
        elif re.match('temp[0-9]*_lcrit', key):
            t_lcrit = value

    if t_in is not None:
        t_dict = {'lcrit'   : t_lcrit,
                  'min'     : t_min,
                  'fan_min' : t_fan_min,
                  'fan_max' : t_fan_max,
                  'max'     : t_max,
                  'crit'    : t_crit,
        }
        msg = 'temp:%2.1f C' %(t_in)
        cnt = 0
        for key, val in t_dict.items():
            if val is not None:
                if cnt == 0:
                    msg += ' ('
                    cnt = 1
                else:
                    msg += ', '
                msg += '%s = %d C' %(key, int(val))
        if cnt:
            msg += ')'
        print(msg)

def print_fan(m):
    f_in = None; f_max = None; f_min = None; f_crit = None; f_var = None; f_dir = None
    f_lcrit = None; f_max_hyst = None

    for key,value in m.items():
        if re.match('fan[0-9]*_input', key):
            f_in = value
        elif re.match('fan[0-9]*_max', key):
            f_max = value
        elif re.match('fan[0-9]*_min', key):
            f_min = value
        elif re.match('fan[0-9]*_var', key):
            f_var = value
        elif re.match('fan[0-9]*_speed_tolerance', key):
            f_var = value
        elif re.match('fan[0-9]*_dir', key):
            if value == 1:
                f_dir = "B2F"
            elif value == 0:
                f_dir = "F2B"
            else:
                f_dir = "none"

    if f_in is not None:
        print('fan:%d RPM   (max = %d RPM, min = %d RPM, limit_variance = %d direction = %s)' \
              %(f_in, f_max, f_min, f_var, f_dir))

def print_volt(m):
    v_in = None; v_max = None; v_min = None; v_crit = None
    v_lcrit = None; v_max_hyst = None

    for key,value in m.items():
        if re.match('in[0-9]*_input', key):
            v_in = value
        elif re.match('in[0-9]*_max', key):
            v_max = value
        elif re.match('in[0-9]*_min', key):
            v_min = value
    if v_in is not None:
        print('voltage:%.2f V   (max = %.2f V, min = %.2f V)' \
              %(v_in, v_max, v_min))

def print_system(m):
    # For system health components, display the state and any additional info
    msg = 'state:%s' % m.get('state', 'UNKNOWN')
    # Add any additional system-specific information if available
    for key, value in m.items():
        if key not in ['name', 'description', 'type', 'state', 'prev_state', 'msg']:
            if value is not None:
                msg += ', %s:%s' % (key, value)
    print(msg)

def print_power(m):
    p_in = None; p_cap_in = None; v_in = []; c_in = []
    msg = ''
    if m['state'] == 'OK':
        for key,value in m.items():
            if re.match('psu[0-9]*_power_out$', key):
                p_in = value
            elif re.match('psu[0-9]*_power_out_max$', key):
                p_cap_in = value
            elif re.match('psu[0-9]*_volt_out$', key):
                v_in.append(value)
            elif re.match('psu[0-9]*_curr_out$', key):
                c_in.append(value)
            """
            elif re.match('in[0-9]*_input$', key):
                v_in.append(value)
            elif re.match('curr[0-9]*_input$', key):
                c_in.append(value)
            """

        if p_in is not None:
            msg = 'power:%.1f W   ' %p_in
        if p_cap_in is not None:
            msg += 'power capacity:%.1f W   ' %p_cap_in
        b = 0
        if len(v_in) > 0:
            f = []
            for x in v_in:
                f.append('%.2f' %x)
            if b == 0:
                msg += '('; b = 1
            msg += 'voltages = %s V' %f
        if len(c_in) > 0:
            f = []
            for x in c_in:
                f.append('%.2f' %x)
            if b == 0:
                msg += '('; b = 1
            msg += ' currents = %s A' %f
        if b:
            msg += ')'
    if msg:
        print(msg)

# check to see if an instance is already running
#
def already_running(pidfile, prog):
    myname=os.path.basename(prog)
    try:
        if not os.path.isfile(pidfile):
            return False
        f = open(pidfile, 'r')
        oldpid = re.findall('\D*(\d+).*', (f.readline()))[0]
        f.close()
        if not os.path.exists('/proc/%s' % oldpid):
            return False
        pfile = open('/proc/%s/cmdline' % oldpid, 'r')
        if myname not in pfile.readline():
            pfile.close()
            return False
        pfile.close()
        return True
    except Exception as inst:
        raise SMONDRuntimeError("unable to validate pidfile %s: %s" %
                                (pidfile, str(inst)))

def sensor_sort(key):
    match = re.match('(.*)(\d+)', key)
    if match:
        return '%s%04u' % (match.group(1), int(match.group(2)))
    return key

def validate_data(value):
    if not value or isinstance(value, (list, int, float)):
        return value
    elif isinstance(value, str):
        try:
            value.decode('utf-8')
            return value
        except UnicodeDecodeError:
            return "ERROR: Invalid data in output"
    return "ERROR: Data type unknown"

def psu_sensor(m):
    if m['name'].startswith('PSU'):
        return True

    return False    

#-------------------------------------------------------------------------------
#
# Main
#
def main():

    pidfile = "/var/run/smond.pid"
    if not already_running(pidfile, 'smond'):
        print('smond is not running. Start smond and then run smonctl')
        sys.exit(-1)

    parser = argparse.ArgumentParser(description="Displays hardware sensors data")
    parser.add_argument('-j', '--json', help ="Generate JSON output",
                        action="store_true")
    parser.add_argument('-s', '--sensor', help ="Display data for the given sensor")
    parser.add_argument('-c', '--check', help ="Checks the status of all the sensors", 
                        action="store_true")
    parser.add_argument('-r', '--relax_check', help ="Checks the status of all the sensors except PSUs", 
                        action="store_true")
    parser.add_argument('-v', '--verbose', help ="Display detailed hardware sensors data",
                        action="store_true")
    parser.add_argument('-t', '--type',
                        help = "Display data for given sensor type",
                        choices = ('temp', 'power', 'fan', 'volt', 'system'))

    args = parser.parse_args()
    ret = 0
    options = { 'temp'  : print_temp,
                'power' : print_power,
                'fan'   : print_fan,
                'volt'  : print_volt,
                'system': print_system,
                }
    files  = os.listdir(path)
    msg = ''
    json_output = []
    not_ok_count = 0

    for f in sorted(files, key=sensor_sort):
        m = unpickle(f)
        if not m:
            continue
        if args.sensor:
            if m['name'].upper() == args.sensor.upper():
                disp = True
            else:
                disp = False
        else:
            disp = True
        if args.type:
            if m['type'] == args.type:
                disp = True
            else:
                disp = False
        if disp == False:
            continue
        if args.json:
            if m['type'] == 'power':
                for key, value in list(m.items()):
                    if key.startswith('power'):
                        m[key.split('_', 1)[1]] = validate_data(m.pop(key))
                        if key.startswith('fan') and key.endswith('input'):
                            if m['state'] == 'ABSENT':
                                # If the fan is absent, report the speed as 0.  
                                m[key.split('_', 1)[1]] = 0
            else:
                for key, value in list(m.items()):
                    if key.startswith('fan') or key.startswith('temp') or key.startswith('in'):
                        m[key.split('_', 1)[1]] = validate_data(m.pop(key))
                        if key.startswith('fan') and key.endswith('input'):
                            if m['state'] == 'ABSENT':
                                # If the fan is absent, report the speed as 0.  
                                m[key.split('_', 1)[1]] = 0
            json_output.append(m)
        elif args.verbose:
            p = '%s' %m['name']
            if m['name'] != m['description']:
                p += '(%s)' %m['description']
            p += ':  %s' %m['state']
            print(p)
            if m['state'] != 'ABSENT':
                options[m['type']](m)
            print('')
            if m.get('msg') is not None:
                msg += m['msg'] + '\n'
        elif args.check or args.relax_check:
            if args.relax_check and psu_sensor(m):
                continue
            if m['state'] != 'OK':
                not_ok_count += 1
        else:
            p = '%-14s' %m['name']
            if m['name'] != m['description']:
                p += '(' + '%-42s' %m['description'] + ')'
            else:
                p += '%44s' %('')
            p += ':  %s' %m['state']
            print(p)

    if args.json:
        print(json.dumps(json_output, indent=4))
    elif args.verbose and msg:
        print('Messages:')
        print(msg)
    elif args.check or args.relax_check:
        if not_ok_count:
            print('FAIL')
        else:
            print('PASS')

    return ret

#--------------------
#
# execution check
#
if __name__ == "__main__":
    try:
        signal.signal(signal.SIGTERM, exit_normally)
        # Cause all warnings to always be triggered.
        warnings.simplefilter("always")
        warnings.formatwarning = smondwarn
        exit(main())
    except SMONDRuntimeError as errstr:
        sys.stderr.write("%s : ERROR : %s\n" % (sys.argv[0], str(errstr)))
        exit(1)
    except KeyboardInterrupt:
        exit_normally()
    except Exception:
        (exc_type, exc_value, exc_traceback) = sys.exc_info()
        err = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
        log = 'Unhandled Exception : %s' % err
        sys.stderr.write(log)
