#!/usr/share/venvs/netq-agent/bin/python
#
# Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited
#


import json
import os
import sys
import time
import random
import datetime
import argparse
import shutil

class QueueHistogramDataGenerator(object):
    snapshot_dir = "{}/{}".format(os.path.expanduser('~'), "histograms")
    monitor_config_file = "/etc/cumulus/datapath/monitor.conf"

    def __init__(self, port_set, interval=1, n_records=1000, tc=0, start_time=None, sample_population_max=1000000, ifprefix="swp"):
        self.port_set = port_set.copy()
        self.interval = interval
        self.n_records = n_records
        self.tc = tc
        self.start_time = self.get_timestamp(start_time) if start_time != None else start_time
        self.sample_population_max = sample_population_max
        self.ifprefix = ifprefix

    def get_timestamp(self, start_time=None):
        stime = 0
        if type(start_time) == str:
            time_format = '%Y-%m-%d %H:%M:%S.%f'
            stime = int(datetime.datetime.strptime(start_time.split('+')[0], time_format).timestamp())
        elif type(start_time) == int:
            stime = start_time
        else:
            stime = int(datetime.datetime.now().timestamp())

        print("Setting timestamp = {} - {}".format(stime, datetime.datetime.fromtimestamp(stime)))

        return stime

    def generate_bin_values(self, n_bins, sample_max=1000000, bin_values=[]):
        for bin in range(0, n_bins):
            bin_values.append(random.randint(1, sample_max - sum(bin_values) - n_bins + bin))
        return bin_values.copy()

    def generate_queue_histogram(self, time_step):
        q_hist_data = {}
        timestamp_str = datetime.datetime.fromtimestamp(time_step).strftime('%Y-%m-%d %H:%M:%S.%f')
        q_hist_config = {"bin_min": 864, "bin_max": 10080, "bin_size": 1152, "sample_time": 1024}
        q_hist_data["version"] = "3.0"
        q_hist_data["timestamp_info"] = {"start_datetime": timestamp_str, "end_datetime": timestamp_str}
        q_hist_data["buffer_info"] = "null"
        q_hist_data["packet_info"] = "null"
        q_hist_data["histogram_info"] = {}
        q_hist_data["histogram_tc_info"] = {"version": "1.0"}
        q_hist_data["histogram_pg_info"] = "null"
        q_hist_data["histogram_counter_info"] = "null"
        q_hist_data["histogram_latency_info"] = "null"
        q_hist_data["bandwidth_info"] = "null"
        for swp_port in port_set:
            q_histogram_range = ["0:863", "864:2015", "2016:3167", "3168:4319", "4320:5471", "5472:6623", "6624:7775", "7776:8927", "8928:10079", "10080:*"]
            q_histogram_data = {"0:863": 0, "864:2015": 0, "2016:3167": 0, "3168:4319": 0, "4320:5471": 0, "5472:6623": 0, "6624:7775": 0, "7776:8927": 0, "8928:10079": 0, "10080:*": 0}
            q_histogram_bins = random.randint(1, 10)
            q_histogram_bin_start_offset = random.randint(0, q_histogram_bins)
            #print("Offset: {}, Bins: {}".format(q_histogram_bin_start_offset, q_histogram_bins))
            if (q_histogram_bin_start_offset + q_histogram_bins) >= 10:
                q_histogram_bin_start_offset -= (q_histogram_bin_start_offset + q_histogram_bins) % 10

            q_histogram_bin_values = self.generate_bin_values(q_histogram_bins, self.sample_population_max, [])
            bin_idx = 0
            for off in range(q_histogram_bin_start_offset, q_histogram_bins):
                key = q_histogram_range[off]
                q_histogram_data[key] = q_histogram_bin_values[bin_idx]
                bin_idx += 1

            q_hist_data["histogram_info"][swp_port] = q_histogram_data
            q_hist_data["histogram_tc_info"][swp_port] = {str(self.tc): {"data": q_histogram_data, "config": q_hist_config}}

        return q_hist_data

    def generate_queue_histograms(self, file_count=200, file_path=snapshot_dir):
        count = 0
        fileidx = 0
        basefilename = file_path + "/histogram_stats_"
        time_step = self.start_time if self.start_time != None else self.get_timestamp()
        while (True):
            if (self.n_records > 0 and count >= self.n_records):
                break
            json_data = json.dumps(self.generate_queue_histogram(time_step))
            filename = basefilename + str(fileidx)
            file = open(filename, "w")
            file.write(json_data)
            file.close()
            print("Timestep - {}".format(datetime.datetime.fromtimestamp(time_step).strftime('%Y-%m-%d %H:%M:%S.%f')))
            print("Wrote file {} @ {}".format(filename, datetime.datetime.now()))
            time.sleep(self.interval)
            count += 1
            fileidx = (fileidx + 1) % file_count
            if self.start_time != None:
                time_step += self.interval
            else:
                time_step = self.get_timestamp()

def do_setup_configuration(config_file, nports, nsplits, tc):
    # Create histogram snapshot path
    snapshot_dir = QueueHistogramDataGenerator.snapshot_dir
    if not os.path.exists(snapshot_dir):
        os.mkdir(snapshot_dir)

    # Copy the existing config file
    bck_config_file = config_file + ".bck"
    if os.path.exists(config_file):
        shutil.copy(config_file, bck_config_file)

    if nsplits == 0:
        ports = "swp{}-swp{}".format(1, nports)
    else:
        port_list = []
        for port in range(1, nports + 1):
            ports = "swp{}s{}-swp{}s{}".format(port, 0, port, nsplits - 1)
            port_list.append(ports)
        ports = ",".join(port_list)
    histogram_config = """
    monitor.port_group_list = [histogram_pg]
    monitor.histogram_pg.port_set            = {ports}
    monitor.histogram_pg.stat_type           = histogram
    monitor.histogram_pg.cos_list            = [{tcval}]
    monitor.histogram_pg.trigger_type        = timer
    monitor.histogram_pg.timer               = 1s
    monitor.histogram_pg.action_list         = [snapshot]
    monitor.histogram_pg.snapshot.file       = {snapshot_dir}/histogram_stats
    monitor.histogram_pg.snapshot.file_count = 200
    monitor.histogram_pg.histogram.minimum_bytes_boundary =   960
    monitor.histogram_pg.histogram.histogram_size_bytes   = 12288
    monitor.histogram_pg.histogram.sample_time_ns         =  1024
    """.format(ports=ports, tcval=tc, snapshot_dir=snapshot_dir)

    conf_file = open(config_file, "w")
    conf_file.write(histogram_config)
    conf_file.close()

def do_cleanup_configuration(config_file):
    snapshot_dir = QueueHistogramDataGenerator.snapshot_dir
    config_file = QueueHistogramDataGenerator.monitor_config_file
    bck_config_file = config_file + ".bck"
    shutil.move(bck_config_file, config_file)
    shutil.rmtree(snapshot_dir, ignore_errors=False)

def setup(args, nports, nsplits, tc):
    port_set = []
    if nsplits == 0:
        for port in range(1, nports + 1):
            port_set.append("swp{}".format(port))
    else:
        for port in range(1, nports + 1):
            for split in range(0, nsplits):
                port_set.append("swp{}s{}".format(port, split))
    if args.setupconf:
        do_setup_configuration(QueueHistogramDataGenerator.monitor_config_file, nports, nsplits, tc)

    return port_set

def cleanup(args):
    if args.setupconf:
        do_cleanup_configuration(QueueHistogramDataGenerator.monitor_config_file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--randseed", help="random seed")
    parser.add_argument("--nports", help="number of ports")
    parser.add_argument("--nsplits", help="number splits per port")
    parser.add_argument("--interval", help="data generation interval")
    parser.add_argument("--nrecords", help="total number of records to generate")
    parser.add_argument("--tc", help="traffic class")
    parser.add_argument("--starttime", help="timestamp of the data")
    parser.add_argument("--filecount", help="file counts")
    parser.add_argument("--filepath", help="filepath where the data would be generated")
    parser.add_argument("--samplesize", help="Sample population size")
    parser.add_argument("--ifprefix", help="Interface prefix")
    parser.add_argument("--setupconf", help="Automatically setup configuration", action="store_true")

    args = parser.parse_args()

    rand_seed = int(datetime.datetime.now().timestamp()) if args.randseed == None else int(args.randseed)

    n_ports = 64 if args.nports == None else int(args.nports)
    n_splits = 2 if args.nports == None else int(args.nsplits)
    tc = random.randint(0, 7) if args.tc == None else int(args.tc)
    interval = 1 if args.interval == None else int(args.interval)
    n_records = 2592000 if args.nrecords == None else int(args.nrecords) # 30 days default
    start_time = args.starttime if args.starttime != None else None
    sample_size = 1000000 if args.samplesize == None else int(args.samplesize)

    file_count = 200 if args.filecount == None else int(args.filecount)
    file_path = QueueHistogramDataGenerator.snapshot_dir if args.filepath == None else args.filepath
    ifprefix = "swp" if args.ifprefix == None else args.ifprefix
    setupconf = True if args.setupconf else False

    port_set = setup(args, n_ports, n_splits, tc)

    random.seed(rand_seed)
    q_hist_data_generator = QueueHistogramDataGenerator(port_set, interval, n_records, tc, start_time, sample_size, ifprefix)

    print("Configuration:")
    print(" Random seed         - {}".format(rand_seed))
    print(" Start time          - {}".format(start_time))
    print(" Number of ports     - {}".format(n_ports))
    print(" Number of splits    - {}".format(n_splits))
    print(" Interval            - {}".format(interval))
    print(" Number of records   - {}".format(n_records))
    print(" Traffic class       - {}".format(tc))
    print(" Sample size         - {}".format(sample_size))
    print(" File count          - {}".format(file_count))
    print(" File path           - {}".format(file_path))
    print(" Interface prefix    - {}".format(ifprefix))
    print(" Setup configuration - {}".format(setupconf))

    print(" Histogram generation started for every {} sec(s)...".format(interval))
    q_hist_data_generator.generate_queue_histograms(file_count=file_count, file_path=file_path)
    print("Histogram generation ended...")

    cleanup(args)
