#!/usr/share/venvs/netq-agent/bin/python
#
# Copyright (c) 2023 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 subprocess
import os
import sys
from time import localtime
from netq_lib.common.utils import subprocess_wrapper

DATE = '%d%d%d' % (localtime().tm_year, localtime().tm_mon, localtime().tm_mday)
TIME = '%d%d%d' % (localtime().tm_hour, localtime().tm_min, localtime().tm_sec)
SVC_PROP = "Names,ActiveState,UnitFileState,MainPID,ActiveEnterTimestamp"
PIP_PATH = "/usr/share/venvs/netq-agent/bin/pip"
PY_SPY_PATH = "/usr/share/venvs/netq-agent/bin/py-spy"


if __name__ == '__main__':
    prog_name = os.path.basename(sys.argv[0])

    if os.getuid():
        print('Must run {} as root'.format(prog_name))
        sys.exit(1)

    if not os.path.exists('/bin/systemctl'):
        print('Non systemctl systems not supported')
        sys.exit(1)

    cmd = '%s show -p %s %s' % ('/bin/systemctl', SVC_PROP, "netq-agent")
    cmd_out = subprocess_wrapper(cmd.split())

    service_state = dict([line.split("=") for line in cmd_out.splitlines()])
    if service_state['ActiveState'] != "active":
        print("NetQ Agent service is not active")
        sys.exit(1)

    if not service_state['MainPID']:
        print("NetQ Agent PID not found. Check if it is running")
        sys.exit(1)
    pid = service_state['MainPID']

    print("installing the py-spy profile tool")
    subprocess_wrapper([PIP_PATH, "install", "py-spy"])

    print("profiling netq-agent(pid: {}) for 2 minute. Please wait..".format(pid))

    if not os.path.exists(PY_SPY_PATH):
        print("py-spy installation failed")
        sys.exit(1)
    PROFILE_OUT_FILE = "/var/log/netq-agent-%s-%s-%s.svg" % (pid, DATE, TIME)
    cmd = "%s record -o %s -p %s -s -d 120" % (PY_SPY_PATH, PROFILE_OUT_FILE, pid)
    try:
        subprocess_wrapper(cmd.split())
        print('Please share the file {} '.format(PROFILE_OUT_FILE))
    except subprocess.CalledProcessError as ex:
        print('ERROR: could not profile agent ({})'.format(ex.cmd))
        pass
