#!/usr/bin/python3
#
# Copyright © 2021 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation and its affiliates
# (the "Company") and all right, title, and interest in and to the software
# product, including all associated intellectual property rights, are and
# shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.

import argparse
import os
import subprocess
import sys
import signal
import re

import pyinotify
from time import sleep
import subprocess
import json
from pydash import py_

COMMANDS = [
    'job-schedule',
    'job-cancel',
    'profile-add',
    'profile-delete',
    'target-add',
    'target-delete'
]

HFT_CONFIG_DIR = '/etc/cumulus/telemetry/hft/'
HFT_CONFIG_FILE = HFT_CONFIG_DIR + 'hft.conf'
TMP_CONFIG_FILE = HFT_CONFIG_DIR + '_hft.conf'
HFT_JOB_CONFIG_FILE = HFT_CONFIG_DIR + 'hft_job.conf'

BACKUP_CONFIG_FILE = HFT_CONFIG_DIR + 'hft.conf.backup'

HFT_STATUS_FILE = '/var/run/cumulus/hft/hft_config.status'
HFT_STATUS_TIMEOUT = 30000

verbose = False
hft_config_changed = False

def verbose_print(arg):
    if verbose:
        print(arg)


def print_help():
    help_str='''
Usage: cl-hft-tool { COMMAND | help } [ OPTIONS ]

COMMAND:
    job-schedule
    job-cancel
    profile-add
    profile-delete
    target-add
    target-delete

OPTIONS:
    --time      Start date and time of format YYYY-MM-DD-HH:MM:SS
    --duration  Duration in seconds
    --profile   profile name
    --ports     Interface list
    --job       job id
    --counter   counter
    --tc        tc list
    --interval  sample interval
    --cf-only   Config file only mode, update only config file , do not trigger service. This will be helpful in batching multiple add/deletes and then triggering service to read profiles

Command to start a job:
	cl-hft-tool  job-schedule -t 2024-05-30-12:00:00 -d 20 -p standard -I swp1-swp64

Command to cancel a job:
	cl-hft-tool  job-cancel -j 1

Command to Add a profile:
	cl-hft-tool  profile-add -p profile1 -c tx-byte -tc 0-3 -si 1000

Command to Delete a profile:
	cl-hft-tool  profile-delete -p profile1
    cl-hft-tool  profile-delete -p all

Command to Add target:
	cl-hft-tool  target-add --target influxdb –host “10.101.1.1” –port 8086 –bucket hft-data –org nvidia –token “*****”

Command to Delete target:
	cl-hft-tool  target-delete --target influxdb
'''
    print(help_str)

def already_running(pidfile):
    myname='asic-monitor'
    try:
        if not os.path.isfile(pidfile):
            return False
        oldpid = re.findall('\D*(\d+).*', (open(pidfile, 'r').readline()))[0]
        if not os.path.exists('/proc/%s' % oldpid):
            return False
        if myname not in open('/proc/%s/cmdline' % oldpid, 'r').readline():
            return False
        verbose_print("%s already running as process %s\n" % (myname, oldpid))
        return True
    except Exception as e:
        print("unable to validate pidfile %s: %s" % (pidfile, str(e)))

def get_pid_of_asic_monitor():
    pid = 0
    pidfile = "/var/run/asic-monitor.pid"
    if already_running(pidfile):
        pid = int(re.findall('\D*(\d+).*', (open(pidfile, 'r').readline()))[0])
        verbose_print("get_pid_of_asic_monitor pid %s\n" %(pid))
    return pid    

def start_asic_monitor():
    try:
        os.system('systemctl reset-failed asic-monitor')
        os.system('systemctl start asic-monitor')
        # sleep here it to allow asic-monitor register for signals
        # if its signaled too quickly then service gets killed
        sleep(5)
        verbose_print("asic monitor started\n")
    except subprocess.CalledProcessError as e:
        verbose_print("asic monitor start failed")


def handle_job_schedule(args):
    verbose_print(f"scheduling job time - {args.time}, duration - {args.duration},"
          f" profile - {args.profile}, ports - {args.ports},"
          f" description - {args.description}")
    desc = 'NA'
    if args.description:
        desc = args.description

    job_config=f'''hft.action_type = schedule
hft.schedule.start_time = {args.time} 
hft.schedule.duration = {args.duration}
hft.schedule.port_set = {args.ports}
hft.schedule.profile_name = {args.profile}
hft.schedule.description = {desc}
'''
    os.makedirs(os.path.dirname(HFT_CONFIG_DIR), exist_ok=True)
    with open(HFT_JOB_CONFIG_FILE, 'w+') as temp_file:
        temp_file.write(f'{job_config}')


def handle_job_cancel(args):
    verbose_print(f"canceling job job id - {args.job}")
    job_config=f'''hft.action_type = cancel
hft.cancel.job_id = {args.job}
'''
    os.makedirs(os.path.dirname(HFT_CONFIG_DIR), exist_ok=True)
    with open(HFT_JOB_CONFIG_FILE, 'w+') as temp_file:
        temp_file.write(f'{job_config}')


def get_current_profile_list(line):
    items = line.split('=')
    list = items[1].strip()
    list = list[1:-1]
    verbose_print(f"current counter list is {list}")
    return list

def add_config_to_list(line, config):
    items = line.split('=')
    config_list = items[1].strip()
    config_list = config_list[1:-1]
    config_list = config_list.split(',')
    config_list = list(set(config_list))
    if config not in config_list:
        config_list.append(config)
    
    # list_str = ','.join(config_list)
    list_str = ''
    for item in config_list:
        if list_str:
            list_str += ','
        list_str += item

    new_config = f'{items[0]}= [{list_str}]\n'
    verbose_print(f"new config {new_config}")
    return new_config

def remove_config_from_list(line, config):
    items = line.split('=')
    config_list = items[1].strip()
    config_list = config_list[1:-1]
    config_list = config_list.split(',')
    config_list = list(set(config_list))
    if config in config_list:
        config_list.remove(config)
    
    list_str = ','.join(config_list)
    new_config = f'{items[0]}= [{list_str}]\n'
    verbose_print(f"new config {new_config}")
    return new_config


def check_current_config(line, config):
    items = line.split('.')
    if len(items) > 1 and items[1] == config:
        return True
    return False

def handle_profile_add(args):
    verbose_print(f"adding profile name - {args.name}, counter - {args.counter},"
          f" tc - {args.tc}, interval - {args.interval}")
    global hft_config_changed
    profile_config = f'''hft.{args.name}.counters_list = [{args.counter}]
hft.{args.name}.sample_interval = {args.interval}
hft.{args.name}.tc_list = [{args.tc}]
'''
    os.makedirs(os.path.dirname(HFT_CONFIG_DIR), exist_ok=True)
    with open(TMP_CONFIG_FILE, 'w+') as temp_file:
        try:
            with open(HFT_CONFIG_FILE, 'r') as curr_file:
                for line in curr_file:
                    if check_current_config(line, args.name):
                        continue
                    if 'profile_list' in line:
                        line = add_config_to_list(line, args.name)
                    temp_file.write(f'{line}')
        except FileNotFoundError:
            profile_list_cfg = f'hft.profile_list = [{args.name}]\n'
            temp_file.write(f'{profile_list_cfg}')
        temp_file.write(f'{profile_config}')

    os.renames(TMP_CONFIG_FILE, HFT_CONFIG_FILE)
    hft_config_changed = True


def handle_profile_delete(args):
    verbose_print(f"deleting profile name - {args.name}")
    global hft_config_changed
    os.makedirs(os.path.dirname(HFT_CONFIG_DIR), exist_ok=True)
    with open(TMP_CONFIG_FILE, 'w+') as temp_file:
        if args.name == 'all':
            # FIXME: this will clean up target configs also
            temp_file.write('hft.profile_list = []\n')
        else:
            try:
                with open(HFT_CONFIG_FILE, 'r') as curr_file:
                    for line in curr_file:
                        if check_current_config(line, args.name):
                            continue
                        if 'profile_list' in line:
                            line = remove_config_from_list(line, args.name)
                        temp_file.write(f'{line}')
            except FileNotFoundError:
                print('hft config file not present')
                return

    os.rename(TMP_CONFIG_FILE, HFT_CONFIG_FILE)
    hft_config_changed = True

def handle_target_add(args):
    verbose_print(f"adding target {args.target}")
    global hft_config_changed
    if args.target == 'local':
        return
    else:
        sys.stderr.write('invalid target, only local target supported')
        exit(1)

#     if args.target == 'influxdb':
#         target_config = f'''
# hft.influxdb.host = {args.host}
# hft.influxdb.port = {args.port}
# hft.influxdb.bucket = {args.bucket}
# hft.influxdb.bucket = {args.org}
# hft.influxdb.token = {args.token}
# '''
#     os.makedirs(os.path.dirname(HFT_CONFIG_DIR), exist_ok=True)
#     with open(TMP_CONFIG_FILE, 'w+') as temp_file:
#         target_added = False
#         try:
#             with open(HFT_CONFIG_FILE, 'r') as curr_file:
#                 for line in curr_file:
#                     if check_current_config(line, args.target):
#                         continue
#                     if 'target' in line:
#                         line = add_config_to_list(line, args.target)
#                         target_added = True
#                     temp_file.write(f'{line}')
#         except FileNotFoundError:
#             target_list = f'hft.target = [{args.target}]\n'
#             temp_file.write(f'{target_list}')
#         if not target_added:
#             target_list = f'hft.target = [{args.target}]\n'
#             temp_file.write(f'{target_list}')
#         temp_file.write(f'{target_config}')

#     os.rename(TMP_CONFIG_FILE, HFT_CONFIG_FILE)
#     hft_config_changed = True


def handle_target_delete(args):
    verbose_print(f"deleting target {args.target}")
    # global hft_config_changed
    # with open(TMP_CONFIG_FILE, 'w+') as temp_file:
    #     try:
    #         with open(HFT_CONFIG_FILE, 'r') as curr_file:
    #             for line in curr_file:
    #                 if check_current_config(line, args.target):
    #                     continue
    #                 if 'target' in line:
    #                     line = remove_config_from_list(line, args.target)
    #                 temp_file.write(f'{line}')
    #     except FileNotFoundError:
    #         print('hft config file not present')
    #         return

    # os.rename(TMP_CONFIG_FILE, HFT_CONFIG_FILE)
    # hft_config_changed = True

def add_job_parser(subparsers):
    parser_sched = subparsers.add_parser('job-schedule')
    parser_sched.add_argument("--time",
                              required=True,
                              help="now (or) Start date and time of format YYYY-MM-DD-HH:MM:SS")
    parser_sched.add_argument("--duration",
                              required=True,
                              default=20,
                              type=int,
                              help="Duration in seconds")
    parser_sched.add_argument("--profile",
                              required=True,
                              help="profile name")
    parser_sched.add_argument("--ports",
                              required=True,
                              default='all',
                              help="Interface list (range(-) or comma(,) separated)")
    parser_sched.add_argument("--description",
                              help="Job description")
    parser_sched.set_defaults(func=handle_job_schedule)

    parser_cancel = subparsers.add_parser('job-cancel')
    parser_cancel.add_argument("--job",
                               required=True,
                               help="job id")
    # parser_cancel.add_argument("--profile",
    #                            help="profile name")
    parser_cancel.set_defaults(func=handle_job_cancel)


def add_profile_parser(subparsers):
    parser_add = subparsers.add_parser('profile-add')
    parser_add.add_argument("--name",
                            required=True,
                            help="profile name")
    parser_add.add_argument("--counter",
                            required=True,
                            help="counter list(Ex:tc_curr_occupancy for tc "+
                            "and if_in_octets, if_out_octets for  ports)")
    parser_add.add_argument("--tc",
                            required=True,
                            help="tc list(range(-) or comma(,) separated)")
    parser_add.add_argument("--interval",
                            required=True,
                            type=int,
                            help="sample interval in micro seconds")
    parser_add.set_defaults(func=handle_profile_add)

    parser_del = subparsers.add_parser('profile-delete')
    parser_del.add_argument("--name",
                            required=True,
                            help="profile name")
    parser_del.set_defaults(func=handle_profile_delete)


def add_target_parser(subparsers):
    parser_add = subparsers.add_parser('target-add')
    parser_add.add_argument("--target",
                            required=True,
                            help="target type")
    # parser_add.add_argument("--host",
    #                         required=True,
    #                         help="influx DB host")
    # parser_add.add_argument("--port",
    #                         required=True,
    #                         type=int,
    #                         help="influx DB port")
    # parser_add.add_argument("--bucket",
    #                         required=True,
    #                         help="influx DB bucket")
    # parser_add.add_argument("--org",
    #                         required=True,
    #                         help="influx DB organization")
    # parser_add.add_argument("--token",
    #                         required=True,
    #                         help="influx DB token")
    parser_add.set_defaults(func=handle_target_add)

    parser_del = subparsers.add_parser('target-delete')
    parser_del.add_argument("--target",
                            required=True,
                            help="target type")
    parser_del.set_defaults(func=handle_target_delete)


def parse_arguments():
    parser = argparse.ArgumentParser(description='cl-hft-tool')

    subparsers = parser.add_subparsers(title='command',
                                       required=True,
                                       description='valid subcommands',
                                       help='additional help')

    add_job_parser(subparsers)
    add_profile_parser(subparsers)
    add_target_parser(subparsers)

    parser.add_argument("--cf-only",
                        default=False,
                        action='store_true',
                        help="Config file only mode, update only config file")
    
    parser.add_argument("--verbose",
                    default=False,
                    action='store_true',
                    help="Verbose log")

    try:
        args = parser.parse_args()
    except SystemExit:
        # print_help()
        exit()
    return args

def backup_config():
    try:
        os.system(f'cp {HFT_CONFIG_FILE} {BACKUP_CONFIG_FILE}')
    except OSError:
        sys.stderr.write('could not take backup of current file')
        exit(1)

def restore_config():
    # no need to restore if config file not changed
    if not hft_config_changed:
        return

    try:
        os.system(f'cp {BACKUP_CONFIG_FILE} {HFT_CONFIG_FILE}')
    except OSError:
        sys.stderr.write('could not restore old config')
        exit(1)

def check_status():
    timeout = HFT_STATUS_TIMEOUT
    time = 0

    while(time < timeout):
        if os.path.isfile(HFT_STATUS_FILE):
            config_status = {}
            try:
                with open(HFT_STATUS_FILE, 'r') as status_file:
                    config_status = json.load(status_file)
            except Exception as e:
                # sys.stderr.write('failed')
                sys.stderr.write('asic-monitor status file read failed')
                return 1
            
            verbose_print('status: %s\n' % (str(config_status)))
            status = py_.get(config_status, 'status', 'failed')
            msg = py_.get(config_status, 'msg', '')
            
            if status == 'success':
                print(f'{msg}')
                return 0
            
            sys.stderr.write(f'{msg}')
            return 1
        else:
            time += 200
            sleep(0.2)

    sys.stderr.write('asic-monitor status file not found')
    return 1


if __name__ == '__main__':

    args = parse_arguments()
    verbose = args.verbose
    backup_config()
    args.func(args)
    if args.cf_only:
        # Only update config file
        exit()
    pid = get_pid_of_asic_monitor()

    if pid == 0:
        start_asic_monitor()

    pid = get_pid_of_asic_monitor()
    if (pid):
        try: 
            os.unlink(HFT_STATUS_FILE) 
        except FileNotFoundError : 
            pass
        verbose_print("sending SIGHUP to asic-monitor pid %d" %(pid))
        os.kill(pid, signal.SIGHUP)
        ret = check_status()
        if ret == 1:
            restore_config()
            exit(1)
        verbose_print("asic-monitor configured successfully")
        exit(0)

