#!python

import argparse
import json
import sys
from natsort import natsorted

from swsscommon.swsscommon import SonicV2Connector
from tabulate import tabulate

VERSION = '1.0'



def get_psu_status_list():
    psu_status_list = []

    db = SonicV2Connector(host="127.0.0.1")
    db.connect(db.STATE_DB)

    # Currently set chassis_num to 1, need to improve it once new platform API is implemented
    chassis_num = 1
    chassis_name = "chassis {}".format(chassis_num)
    num_psus = db.get(db.STATE_DB, 'CHASSIS_INFO|{}'.format(chassis_name), 'psu_num')
    if not num_psus:
        print('Error: Failed to get the number of PSUs')
        return None

    for psu_idx, psu_key in enumerate(natsorted(db.keys(db.STATE_DB, "PSU_INFO|*")), start = 1):
        psu_key = psu_key.replace("PSU_INFO|", "")
        psu_status = {}

        psu_status['index'] = str(psu_idx)

        psu_name = psu_key
        psu_status['name'] = psu_name

        presence = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'presence')
        psu_status['presence'] = presence

        if presence == 'true':
            oper_status = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'status')
            if oper_status == 'true':
                power_overload = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'power_overload')
                status = 'WARNING' if power_overload == 'True' else 'OK'
            else:
                status = "NOT OK"
        else:
            status = 'NOT PRESENT'
        psu_status['status'] = status

        psu_status['led_status'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'led_status')

        psu_status['model'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'model') if presence else 'N/A'
        psu_status['serial'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'serial') if presence else 'N/A'
        psu_status['revision'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'revision') if presence else 'N/A'
        psu_status['voltage'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'voltage') if presence else 'N/A'
        psu_status['current'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'current') if presence else 'N/A'
        psu_status['power'] = db.get(db.STATE_DB, 'PSU_INFO|{}'.format(psu_name), 'power') if presence else 'N/A'

        psu_status_list.append(psu_status)

    return psu_status_list


def psu_status_show_table(index):
    psu_status_list = get_psu_status_list()

    if not psu_status_list:
        print('Error: Failed to get PSU status')
        return None

    header = ['PSU',  'Model', 'Serial', 'HW Rev', 'Voltage (V)', 'Current (A)', 'Power (W)', 'Status', 'LED']
    status_table = []

    if index > 0:
        if index > len(psu_status_list):
            print("Error: PSU {} is not available. Number of supported PSUs: {}".format(index, len(psu_status_list)))
            return -1

        # Trim the list down to contain only the requested PSU
        psu_status_list = [psu_status_list[index-1]]

    for psu_status in psu_status_list:
        status_table.append([psu_status['name'],
                             psu_status['model'],
                             psu_status['serial'],
                             psu_status['revision'],
                             psu_status['voltage'],
                             psu_status['current'],
                             psu_status['power'],
                             psu_status['status'],
                             psu_status['led_status']])

    if status_table:
        print(tabulate(status_table, header, tablefmt="simple", floatfmt='.2f'))

    return 0


def psu_status_show_json(index):
    psu_status_list = get_psu_status_list()

    if not psu_status_list:
        print('Error: Failed to get PSU status')
        return None

    if index > 0:
        if index > len(psu_status_list):
            print("Error: PSU {} is not available. Number of supported PSUs: {}".format(index, len(psu_status_list)))
            return -1

        # Trim the list down to contain only the requested PSU
        psu_status_list = [psu_status_list[index-1]]

    print(json.dumps(psu_status_list, indent=4))
    return 0


def main():
    parser = argparse.ArgumentParser(description='Display the psu status information',
                                     formatter_class=argparse.RawTextHelpFormatter,
                                     epilog="""
Examples:
  psushow -s
  psushow -s -j
  psushow -s -i 1
""")

    parser.add_argument('-s', '--status', action='store_true', help='Show PSU status information')
    parser.add_argument('-i', '--index', type=int, default=-1, help='The index of PSU to display')
    parser.add_argument('-j', '--json', action='store_true', help='Display output in JSON format')
    parser.add_argument('-v', '--version', action='version', version='%(prog)s {}'.format(VERSION))
    args = parser.parse_args()

    status_show = args.status
    psu_index = args.index
    output_json = args.json

    if status_show:
        if output_json:
            ret = psu_status_show_json(psu_index)
        else:
            ret = psu_status_show_table(psu_index)

        if ret != 0:
            print("Error: failed to get PSU status from state DB")
            return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())
